import dataclasses
import enum
from typing import Optional, TYPE_CHECKING

from spimdisasm.common import SortedDict

from ..symbols import Symbol, label_types, get_all_symbols
from .. import log

from .parent_segment_info import ParentSegmentInfo

if TYPE_CHECKING:
    from ...segtypes.common.segment import Segment


class SegmentKind(enum.Enum):
    """
    The type of segment metadata.
    """

    Global = 0
    """
    Globally vissible segment, usually the first segments.

    It has no overlapping addresses with any other segment.
    """

    Overlay = 1
    """
    Segments that have overlapping Vram addresses with other segments.

    This kind of segment have an `exclusive_ram_id` to group them.
    """

    Unknown = 2
    """
    Junk.

    Referenced addresses that we couldn't find to which segment they belong to.
    """

    Absolute = 3
    """
    User-declared symbols that are not part of the address space of any segment.
    """


@dataclasses.dataclass
class SegmentMetadata:
    """
    Metadata for a top-level segment.
    """

    kind: SegmentKind
    name: str

    rom_start: int
    rom_end: int
    vram_start: int
    vram_end: int

    prioritized_segments: list[str]
    """
    Other segments this segment is allowed to see and prioritize.
    """

    exclusive_ram_id: Optional[str]

    segment: Optional["Segment"]
    """
    Reference to the corresponding top-level segment.

    This is `None` for `Unknown` and `Absolute` segments.
    """

    symbols: SortedDict[Symbol]
    """
    All symbols defined in this segment.

    This only contains symbols that are part of this segment.
    Sorted by Vram address.
    """

    def __init__(
        self,
        kind: SegmentKind,
        name: str,
        rom_start: int,
        rom_end: int,
        vram_start: int,
        vram_end: int,
        prioritized_segments: list[str],
        exclusive_ram_id: Optional[str],
        segment: Optional["Segment"],
    ) -> None:
        if rom_start > rom_end:
            log.error(
                f"Error creating segment '{name}': rom_start (0x{rom_start:X}) is larger than rom_end (0x{rom_end:X})"
            )
        if vram_start > vram_end:
            log.error(
                f"Error creating segment '{name}': rom_start (0x{vram_start:X}) is larger than rom_end (0x{vram_end:X})"
            )

        self.kind = kind
        self.name = name
        self.rom_start = rom_start
        self.rom_end = rom_end
        self.vram_start = vram_start
        self.vram_end = vram_end
        self.prioritized_segments = prioritized_segments
        self.exclusive_ram_id = exclusive_ram_id
        self.segment = segment

        self.symbols = SortedDict()

    def in_rom_range(self, rom: int) -> bool:
        if rom < self.rom_start:
            return False
        if rom >= self.rom_end:
            return False
        return True

    def in_vram_range(self, vram: int) -> bool:
        if vram < self.vram_start:
            return False
        if vram >= self.vram_end:
            return False
        return True

    def rom_from_vram(self, vram: int) -> Optional[int]:
        """
        Converts a Vram address into a Rom address.

        Returns `None` if either:
        - This segment is either the Absolute or Unknown special segments.
        - The given Vram is outside the Vram space of this segment.
        """

        if self.kind == SegmentKind.Unknown or self.kind == SegmentKind.Absolute:
            return None

        if not self.in_vram_range(vram):
            return None
        rom = vram - self.vram_start + self.rom_start
        return rom

    def is_owned_segment(self, info: ParentSegmentInfo) -> bool:
        """
        Checks if the passed info corresponds to this segment metadata.
        """
        if self.kind == SegmentKind.Absolute or self.kind == SegmentKind.Unknown:
            return False

        if self.exclusive_ram_id is None and info.exclusive_ram_id is None:
            # Global segment
            if not self.in_vram_range(info.segment_vram):
                return False
            return True

        else:
            # Overlay
            if self.exclusive_ram_id != info.exclusive_ram_id:
                return False
            if self.rom_start != info.segment_rom:
                return False
            return True

    def create_symbol(self, vram: int, allow_addend: bool) -> Symbol:
        """
        Creates an autogenerated symbol (not user-declared).
        """

        if not self.in_vram_range(vram):
            log.write(
                f"\nWARNING: Bug! Adding symbol 0x{vram:08X} to segment '{self.name}' ({self.kind}),\n"
                f"  but the address of the symbol is outside the segment vram range (0x{self.vram_start:08X} ~ 0x{self.vram_end:08X})\n",
                status="warn",
            )

        symbol = self.find_symbol(vram, allow_addend)
        if symbol is None:
            # Pass as much info as we can to this symbol when creating it.
            symbol = Symbol(
                vram,
                rom=self.rom_from_vram(vram),
                segment=self.segment,
                seg_meta=self,
            )
            self.symbols[vram] = symbol
            get_all_symbols().append(symbol)
            if self.kind == SegmentKind.Unknown:
                symbol.unknown_segment = True

        return symbol

    def add_user_symbol(self, sym: Symbol) -> None:
        """
        Add an user declared symbol.
        """

        if not self.in_vram_range(sym.vram_start):
            log.write(
                f"\nWARNING: Adding user-declared symbol '{sym.name}' (0x{sym.vram_start:08X}) to segment '{self.name}' ({self.kind}),\n"
                f"but the address of the symbol is outside the segment vram range (0x{self.vram_start:08X} ~ 0x{self.vram_end:08X})\n"
                f"This is accepted for now, but may become a hard error in the future. Expect buggy symbol references.\n",
                status="warn",
            )

        existing_sym = self.find_symbol(sym.vram_start, True)
        if existing_sym is not None:
            if existing_sym is sym:
                # there's a bug somewhere...
                pass
            elif sym.type in {"alabel", *label_types} or existing_sym in {
                "alabel",
                *label_types,
            }:
                # It is expected for labels to overlap with functions
                pass
            else:
                if existing_sym.given_size is None:
                    existing_size = "None"
                else:
                    existing_size = f"0x{existing_sym.given_size:X}"
                size = f"0x{sym.given_size:X}" if sym.given_size is not None else "None"
                msg = (
                    f"The user declared symbol '{sym.name}' (Vram 0x{sym.vram_start:08X}, size {size}, segment '{sym.segment}')\n"
                    f"  overlaps with the previously defined '{existing_sym.name}' (Vram 0x{existing_sym.vram_start:08X}, size {existing_size}, segment '{existing_sym.segment}')"
                )
                # TODO: Change this into a hard error.
                if False:
                    log.error(f"\nERROR: {msg}")
                else:
                    log.write(f"\nWARNING: {msg}", status="warn")

        self.symbols[sym.vram_start] = sym
        if sym.rom is None:
            sym.rom = self.rom_from_vram(sym.vram_start)
        sym.seg_meta = self

    def find_symbol(self, vram: int, allow_addend: bool) -> Optional[Symbol]:
        """
        Lookup for a symbol given its address,
        optionally checking its size as a ranged search.
        """

        if allow_addend:
            pair = self.symbols.getKeyRight(vram, True)
            if pair is None:
                return None

            symbol_vram, sym = pair
            size = sym.given_size or 1
            if vram >= symbol_vram + size:
                return None
            return sym

        return self.symbols.get(vram)

    def get_prioritized_segments(self) -> list[str]:
        return self.prioritized_segments
