from dataclasses import dataclass
import re
from typing import Dict, List, Optional, Set, TYPE_CHECKING

import spimdisasm

from ..disassembler import disassembler_instance
from pathlib import Path

# circular import
if TYPE_CHECKING:
    from ..segtypes.segment import Segment
    from .metadata.segment_manager import SegmentManager
    from ..util.metadata.segment_metadata import SegmentMetadata

from . import log, options, progress_bar

all_symbols: List["Symbol"] = []
ignored_addresses: Set[int] = set()
to_mark_as_defined: Set[str] = set()

# Initialize a spimdisasm context, used to store symbols and functions
spim_context = spimdisasm.common.Context()

TRUEY_VALS = ["true", "on", "yes", "y"]
FALSEY_VALS = ["false", "off", "no", "n"]

splat_sym_types = {"func", "jtbl", "jtbl_label", "label", "alabel"}
label_types = {"jtbl_label", "label"}

ILLEGAL_FILENAME_CHARS = ["<", ">", ":", '"', "/", "\\", "|", "?", "*"]


def check_valid_type(typename: str) -> bool:
    if typename[0].isupper():
        return True

    if typename in splat_sym_types:
        return True

    if typename in disassembler_instance.get_instance().known_types():
        return True

    return False


def is_truey(str: str) -> bool:
    return str.lower() in TRUEY_VALS


def is_falsey(str: str) -> bool:
    return str.lower() in FALSEY_VALS


def add_symbol(sym: "Symbol", all_symbols_dict: Dict[int, List["Symbol"]]):
    all_symbols.append(sym)
    items = all_symbols_dict.get(sym.vram_start)
    if items is None:
        all_symbols_dict[sym.vram_start] = [sym]
    else:
        items.append(sym)


def to_cname(symbol_name: str) -> str:
    symbol_name = re.sub(r"[^0-9a-zA-Z_]", "_", symbol_name)

    if symbol_name[0] in ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"]:
        symbol_name = "_" + symbol_name

    return symbol_name


def handle_sym_addrs(
    path: Path,
    sym_addrs_lines: List[str],
    all_segments: "List[Segment]",
):
    def get_seg_for_name(name: str) -> Optional["Segment"]:
        for segment in all_segments:
            if segment.name == name:
                return segment
        return None

    def get_seg_for_rom(rom: int) -> Optional["Segment"]:
        for segment in all_segments:
            if segment.contains_rom(rom):
                return segment
        return None

    all_symbols_dict: Dict[int, List["Symbol"]] = {}

    seen_symbols: Dict[str, "Symbol"] = {}
    prog_bar = progress_bar.get_progress_bar(sym_addrs_lines)
    prog_bar.set_description(f"Loading symbols ({path.stem})")
    line: str
    for line_num, line in enumerate(prog_bar):
        line = line.strip()
        if line != "" and not line.startswith("//"):
            comment_loc = line.find("//")
            line_main = line
            line_ext = ""

            if comment_loc != -1:
                line_ext = line[comment_loc + 2 :].strip()
                line_main = line[:comment_loc].strip()

            try:
                assert line.count(";") == 1, "Line must contain a single semi-colon"
                line_split = line_main.split("=")
                name = line_split[0].strip()
                addr = int(line_split[1].strip()[:-1], 0)
            except:
                log.parsing_error_preamble(path, line_num, line)
                log.write("Line must be of the form")
                log.write("<function_name> = <address>; // attr0:val0 attr1:val1 [...]")
                log.write("with <address> in hex preceded by 0x, or dec")
                log.write("")
                raise

            sym = Symbol(addr, given_name=name)

            ignore_sym = False
            if line_ext:
                for info in line_ext.split(" "):
                    if ":" in info:
                        if info.count(":") > 1:
                            log.parsing_error_preamble(path, line_num, line)
                            log.write(f"Too many ':'s in '{info}'")
                            log.error("")

                        attr_name, attr_val = info.split(":")
                        if attr_name == "":
                            log.parsing_error_preamble(path, line_num, line)
                            log.write(
                                f"Missing attribute name in '{info}', is there extra whitespace?"
                            )
                            log.error("")
                        if attr_val == "":
                            log.parsing_error_preamble(path, line_num, line)
                            log.write(
                                f"Missing attribute value in '{info}', is there extra whitespace?"
                            )
                            log.error("")

                        # Non-Boolean attributes
                        try:
                            if attr_name == "type":
                                if not check_valid_type(attr_val):
                                    log.parsing_error_preamble(path, line_num, line)
                                    log.write(
                                        f"Unrecognized symbol type in '{info}', it should be one of"
                                    )
                                    log.write(
                                        [
                                            *splat_sym_types,
                                            *spimdisasm.common.gKnownTypes,
                                        ]
                                    )
                                    log.write(
                                        "You may use a custom type that starts with a capital letter"
                                    )
                                    log.error("")
                                type = attr_val
                                sym.type = type
                                continue
                            if attr_name == "size":
                                size = int(attr_val, 0)
                                sym.given_size = size
                                continue
                            if attr_name == "rom":
                                rom_addr = int(attr_val, 0)
                                sym.rom = rom_addr
                                continue
                            if attr_name == "segment":
                                seg = get_seg_for_name(attr_val)
                                if seg is None:
                                    log.parsing_error_preamble(path, line_num, line)
                                    log.write(f"Cannot find segment '{attr_val}'")
                                    log.error("")
                                else:
                                    # Add segment to symbol
                                    sym.segment = seg
                                continue
                            if attr_name == "name_end":
                                sym.given_name_end = attr_val
                                continue
                            if attr_name == "filename":
                                sym.given_filename = attr_val
                                continue
                            if attr_name == "visibility":
                                sym.given_visibility = attr_val
                                continue
                            if attr_name == "function_owner":
                                sym.function_owner = attr_val
                                continue
                            if attr_name == "align":
                                align = int(attr_val, 0)

                                if align < 0:
                                    log.parsing_error_preamble(path, line_num, line)
                                    log.error(
                                        f"The given alignment for '{sym.name}' (0x{sym.vram_start:08X}) is negative."
                                    )
                                align_shift = (align & (-align)).bit_length() - 1
                                if (1 << align_shift) != align:
                                    log.parsing_error_preamble(path, line_num, line)
                                    log.error(
                                        f"The given alignment '0x{align:X}' for symbol '{sym.name}' (0x{sym.vram_start:08X}) is not a power of two."
                                    )
                                if sym.vram_start % align != 0:
                                    log.parsing_error_preamble(path, line_num, line)
                                    log.error(
                                        f"The symbol '{sym.name}' (0x{sym.vram_start:08X}) is not aligned already to the given alignment '0x{align:X}'."
                                    )

                                sym.given_align = align
                                continue
                        except:
                            log.parsing_error_preamble(path, line_num, line)
                            log.write(
                                f"value of attribute '{attr_name}' could not be read:"
                            )
                            log.write("")
                            raise

                        # Boolean attributes
                        tf_val = (
                            True
                            if is_truey(attr_val)
                            else False
                            if is_falsey(attr_val)
                            else None
                        )
                        if tf_val is None:
                            log.parsing_error_preamble(path, line_num, line)
                            log.write(
                                f"Invalid Boolean value '{attr_val}' for attribute '{attr_name}', should be one of"
                            )
                            log.write([*TRUEY_VALS, *FALSEY_VALS])
                            log.error("")
                        else:
                            if attr_name == "defined":
                                sym.defined = tf_val
                                continue
                            if attr_name == "extract":
                                sym.extract = tf_val
                                continue
                            if attr_name == "ignore":
                                ignore_sym = tf_val
                                continue
                            if attr_name == "force_migration":
                                sym.force_migration = tf_val
                                continue
                            if attr_name == "force_not_migration":
                                sym.force_not_migration = tf_val
                                continue
                            if attr_name == "allow_addend":
                                sym.allow_addend = tf_val
                                continue
                            if attr_name == "dont_allow_addend":
                                sym.dont_allow_addend = tf_val
                                continue
                            if attr_name == "can_reference":
                                sym.can_reference = tf_val
                                continue
                            if attr_name == "can_be_referenced":
                                sym.can_be_referenced = tf_val
                                continue
                            if attr_name == "allow_duplicated":
                                sym.allow_duplicated = True
                                continue
                            if attr_name == "use_non_matching_label":
                                sym.use_non_matching_label = tf_val
                            if attr_name == "absolute":
                                sym.absolute = tf_val

            if ignore_sym:
                if sym.given_size is None or sym.given_size == 0:
                    ignored_addresses.add(sym.vram_start)
                else:
                    spim_context.addBannedSymbolRangeBySize(
                        sym.vram_start, sym.given_size
                    )
                continue

            if sym.absolute:
                sym.defined = True

                # Most attributes make no sense for a absolute symbol.
                # For now just warn about these two.
                if sym.rom is not None:
                    log.parsing_error_preamble(path, line_num, line)
                    log.write(
                        f"Warning: The symbol '{sym.name}' (0x{sym.vram_start:08X}) has both a rom address (0x{sym.rom:08X}) and absolute:True.",
                        status="warn",
                    )
                if sym.segment is not None:
                    log.parsing_error_preamble(path, line_num, line)
                    log.write(
                        f"Warning: The symbol '{sym.name}' (0x{sym.vram_start:08X}) has both an associated segment ({sym.segment}) and absolute:True.",
                        status="warn",
                    )
            else:
                if sym.segment is None and sym.rom is not None:
                    sym.segment = get_seg_for_rom(sym.rom)

                if sym.segment:
                    if (
                        sym.segment.vram_start is not None
                        and sym.segment.vram_end is not None
                        and not sym.segment.contains_vram(sym.vram_start)
                    ):
                        log.parsing_error_preamble(path, line_num, line)
                        log.write(
                            f"Warning: User-declared symbol '{sym.name}' was associated to segment '{sym.segment.name}', "
                            "but its address is outside the segment's vram range.\n"
                            f"  The symbol's Vram 0x{sym.vram_start:08X} is outside from the segment's Vram range 0x{sym.segment.vram_start:08X} ~ 0x{sym.segment.vram_end:08X}",
                            status="warn",
                        )
                    if sym.rom is not None:
                        expected_rom = sym.segment.ram_to_rom(sym.vram_start)
                        if expected_rom is not None and expected_rom != sym.rom:
                            log.parsing_error_preamble(path, line_num, line)
                            log.write(
                                f"Warning: User-declared symbol '{sym.name}' has a wrong user-declared rom address.\n"
                                f"  This symbol has been mapped to segment '{sym.segment.name}', but the expected Rom\n"
                                f"  address for a symbol with Vram address 0x{sym.vram_start:08X} in that segment is\n"
                                f"  0x{expected_rom:X}, but the given Rom address is 0x{sym.rom:X}.",
                                status="warn",
                            )

                sym.user_declared = True

                if sym.name in seen_symbols:
                    item = seen_symbols[sym.name]
                    if not sym.allow_duplicated or not item.allow_duplicated:
                        log.parsing_error_preamble(path, line_num, line)
                        log.error(
                            f"Duplicate symbol detected! {sym.name} has already been defined at vram 0x{item.vram_start:08X}"
                        )

                items = all_symbols_dict.get(addr)
                if items is not None:
                    for item in items:
                        have_same_rom_addresses = sym.rom == item.rom
                        same_segment = sym.segment == item.segment

                        if have_same_rom_addresses and same_segment:
                            if not sym.allow_duplicated or not item.allow_duplicated:
                                log.parsing_error_preamble(path, line_num, line)
                                log.error(
                                    f"Duplicate symbol detected! {sym.name} clashes with {item.name} defined at vram 0x{addr:08X}.\n  If this is intended, specify either a segment or a rom address for this symbol"
                                )

                if len(sym.filename) > 253 or any(
                    c in ILLEGAL_FILENAME_CHARS for c in sym.filename
                ):
                    log.parsing_error_preamble(path, line_num, line)
                    log.error(
                        # sym.name is written on its own line so reading the error message is nicer because the sym name will be very long.
                        # Other lines have two spaces to make identation nicer and consistent
                        f"Ilegal symbol filename detected!\n"
                        f"  The symbol\n"
                        f"    {sym.name}\n"
                        f"  exceeds the 255 bytes filename limit that most OS imposes or uses illegal characters,\n"
                        f"  which will be a problem when writing the symbol to its own file.\n"
                        f"  To fix this specify a `filename` for this symbol, like `filename:func_{sym.vram_start:08X}`.\n"
                        f"  Make sure the filename does not exceed 253 bytes nor it contains any of the following characters:\n"
                        f"    {ILLEGAL_FILENAME_CHARS}"
                    )

                seen_symbols[sym.name] = sym

            add_symbol(sym, all_symbols_dict)


def initialize(all_segments: "List[Segment]"):
    global all_symbols

    all_symbols = []

    # Manual list of func name / addrs
    for path in options.opts.symbol_addrs_paths:
        if path.exists():
            with open(path, encoding="utf-8") as f:
                sym_addrs_lines = f.readlines()
                handle_sym_addrs(path, sym_addrs_lines, all_segments)


def initialize_spim_context(manager: "SegmentManager") -> None:
    spim_context.bannedSymbols |= ignored_addresses

    # Initialize spimdisasm's global ranges.
    if (
        manager.global_rom_start is not None
        and manager.global_rom_end is not None
        and manager.global_vram_start is not None
        and manager.global_vram_end is not None
    ):
        spim_context.changeGlobalSegmentRanges(
            manager.global_rom_start,
            manager.global_rom_end,
            manager.global_vram_start,
            manager.global_vram_end,
        )

    # Pass absolute symbols.
    for sym in manager.absolute_segment.symbols.values():
        add_symbol_to_spimdisasm_segment(spim_context.absoluteSegment, sym)

    all_segment_names = {seg.name for seg in manager.global_segments} | {
        seg.name
        for overlay_cat in manager.overlay_segments.values()
        for seg in overlay_cat.segments.values()
    }

    # Pass global symbols
    for seg_meta in manager.global_segments:
        for sym in seg_meta.symbols.values():
            for prioritized_seg in seg_meta.prioritized_segments:
                if prioritized_seg not in all_segment_names:
                    log.error(
                        f"\nError: The segment '{seg_meta.name}' references non-existing '{prioritized_seg}' segment. Stopping."
                    )
                spim_context.globalSegment.addPrioritizedSegment(prioritized_seg)
            add_symbol_to_spimdisasm_segment(spim_context.globalSegment, sym)

    # Create overlays and pass their symbols.
    for ovl_id, segments_per_rom in manager.overlay_segments.items():
        for seg_meta in segments_per_rom.segments.values():
            spimdisasm_segment = spim_context.addOverlaySegment(
                ovl_id,
                seg_meta.rom_start,
                seg_meta.rom_end,
                seg_meta.vram_start,
                seg_meta.vram_end,
                seg_meta.name,
            )
            for prioritized_seg in seg_meta.prioritized_segments:
                if prioritized_seg not in all_segment_names:
                    log.error(
                        f"\nError: The segment '{seg_meta.name}' references non-existing '{prioritized_seg}' segment. Stopping."
                    )
                spimdisasm_segment.addPrioritizedSegment(prioritized_seg)

            for sym in seg_meta.symbols.values():
                add_symbol_to_spimdisasm_segment(spimdisasm_segment, sym)


def add_symbol_to_spimdisasm_segment(
    segment: spimdisasm.common.SymbolsSegment,
    sym: "Symbol",
) -> spimdisasm.common.ContextSymbol:
    if sym.type == "func" or sym.type == "alabel":
        context_sym = segment.addFunction(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    elif sym.type == "jtbl":
        context_sym = segment.addJumpTable(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    elif sym.type == "jtbl_label":
        context_sym = segment.addJumpTableLabel(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    elif sym.type == "label":
        context_sym = segment.addBranchLabel(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
    else:
        context_sym = segment.addSymbol(
            sym.vram_start, isAutogenerated=not sym.user_declared, vromAddress=sym.rom
        )
        if sym.type is not None:
            context_sym.type = sym.type

    if sym.user_declared:
        context_sym.isUserDeclared = True
    if sym.defined:
        context_sym.isDefined = True
    if sym.rom is not None:
        context_sym.vromAddress = sym.rom
    if sym.given_size is not None:
        context_sym.size = sym.size
    if sym.force_migration:
        context_sym.forceMigration = True
    if sym.force_not_migration:
        context_sym.forceNotMigration = True
    context_sym.functionOwnerForMigration = sym.function_owner
    if sym.allow_addend:
        context_sym.allowedToReferenceAddends = True
    if sym.dont_allow_addend:
        context_sym.notAllowedToReferenceAddends = True
    if sym.can_reference is not None:
        context_sym.allowedToReferenceSymbols = sym.can_reference
    if sym.can_be_referenced is not None:
        context_sym.allowedToBeReferenced = sym.can_be_referenced
    context_sym.setNameGetCallbackIfUnset(lambda _: sym.name)
    if sym.given_name_end:
        context_sym.nameEnd = sym.given_name_end
    if sym.given_visibility:
        context_sym.visibility = sym.given_visibility
    if sym.given_align:
        context_sym.setAlignment(sym.given_align)
    if sym.use_non_matching_label is not None:
        context_sym.useNonMatchingLabel = sym.use_non_matching_label

    return context_sym


# force_in_segment=True when the symbol belongs to this specific parent segment.
# force_in_segment=False when this symbol is just a reference.
def create_symbol_from_spim_symbol(
    most_parent: "Segment",
    current_segment: "Segment",
    context_sym: spimdisasm.common.ContextSymbol,
    *,
    force_in_segment: bool,
) -> "Symbol":
    sym_type = None
    if context_sym.type == spimdisasm.common.SymbolSpecialType.jumptable:
        sym_type = "jtbl"
    elif context_sym.type == spimdisasm.common.SymbolSpecialType.function:
        sym_type = "func"
    elif context_sym.type == spimdisasm.common.SymbolSpecialType.branchlabel:
        sym_type = "label"
    elif context_sym.type == spimdisasm.common.SymbolSpecialType.jumptablelabel:
        sym_type = "jtbl_label"

    in_segment = False
    if (
        context_sym.overlayCategory is None
        and most_parent.get_exclusive_ram_id() is None
    ):
        in_segment = most_parent.contains_vram(context_sym.vram)
    elif context_sym.overlayCategory == most_parent.get_exclusive_ram_id():
        if context_sym.vromAddress is not None:
            in_segment = most_parent.contains_rom(context_sym.vromAddress)
        else:
            in_segment = most_parent.contains_vram(context_sym.vram)

    sym = most_parent.create_symbol(
        context_sym.vram,
        force_in_segment or in_segment,
        type=sym_type,
        reference=True,
    )

    # Avoid overriding names for user declared symbols
    if sym.given_name is not None and context_sym.name is None:
        context_sym.name = sym.given_name
    if sym.given_name is None and context_sym.name is not None:
        sym.given_name = context_sym.name

    # To keep the symbol name in sync between splat and spimdisasm
    if in_segment:
        context_sym.setNameGetCallback(lambda _: sym.name)
    else:
        context_sym.setNameGetCallbackIfUnset(lambda _: sym.name)

    if context_sym.size is not None:
        sym.given_size = context_sym.getSize()
    if context_sym.vromAddress is not None:
        sym.rom = context_sym.getVrom()
    if context_sym.isDefined:
        sym.defined = True
    if context_sym.referenceCounter > 0:
        sym.referenced = True

    # Void the autodetected symbol type if it is a branch target, but the
    # symbol isn't part of a text section.
    # This may happen on handwritten asm where the function jumps to some
    # data/bss symbol. This can be seen on libultra's monoutil.s (__isExp).
    if (
        not current_segment.is_text()
        and sym_type is None
        and sym.type in ("label", "jtbl_label")
    ):
        sym.type = None

    return sym


def mark_c_funcs_as_defined():
    for symbol in all_symbols:
        if len(to_mark_as_defined) == 0:
            return
        sym_name = symbol.name
        if sym_name in to_mark_as_defined:
            symbol.defined = True
            to_mark_as_defined.remove(sym_name)


@dataclass
class Symbol:
    vram_start: int

    given_name: Optional[str] = None
    given_name_end: Optional[str] = None
    rom: Optional[int] = None
    type: Optional[str] = None
    given_size: Optional[int] = None
    segment: Optional["Segment"] = None

    defined: bool = False
    referenced: bool = False
    extract: bool = True
    user_declared: bool = False

    force_migration: bool = False
    force_not_migration: bool = False
    function_owner: Optional[str] = None

    allow_addend: bool = False
    dont_allow_addend: bool = False

    can_reference: Optional[bool] = None
    can_be_referenced: Optional[bool] = None

    linker_section: Optional[str] = None

    allow_duplicated: bool = False

    given_filename: Optional[str] = None
    given_visibility: Optional[str] = None

    given_align: Optional[int] = None

    use_non_matching_label: Optional[bool] = None

    unknown_segment: bool = False
    absolute: bool = False
    seg_meta: Optional["SegmentMetadata"] = None

    _generated_default_name: Optional[str] = None
    _last_type: Optional[str] = None

    def __str__(self):
        return self.name

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Symbol):
            return False
        return self.vram_start == other.vram_start and self.segment == other.segment

    # https://stackoverflow.com/a/56915493/6292472
    def __hash__(self):
        return hash((self.vram_start, self.segment))

    def format_name(self, format: str) -> str:
        ret = format

        ret = ret.replace("$VRAM", f"{self.vram_start:08X}")

        if "$ROM" in ret:
            if not isinstance(self.rom, int):
                log.error(
                    f"Attempting to rom-name a symbol with no ROM address: {self.vram_start:08X} typed {self.type}"
                )
            ret = ret.replace("$ROM", f"{self.rom:X}")

        if "$SEG" in ret:
            if self.segment is None:
                # This probably is fine - we can't expect every symbol to have a segment. Fall back to just the ram address
                return f"{self.vram_start:X}"
            assert self.segment is not None
            ret = ret.replace("$SEG", self.segment.name)

        return ret

    @property
    def default_name(self) -> str:
        if self._generated_default_name is not None:
            if self.type == self._last_type:
                return self._generated_default_name

        if self.segment:
            if isinstance(self.rom, int):
                suffix = self.format_name(self.segment.symbol_name_format)
            else:
                suffix = self.format_name(self.segment.symbol_name_format_no_rom)
        else:
            if isinstance(self.rom, int):
                suffix = self.format_name(options.opts.symbol_name_format)
            else:
                suffix = self.format_name(options.opts.symbol_name_format_no_rom)

        if self.type == "func":
            prefix = "func"
        elif self.type == "jtbl":
            prefix = "jtbl"
        elif self.type in {"jtbl_label", "label"}:
            return f".L{suffix}"
        else:
            prefix = "D"

        self._last_type = self.type
        self._generated_default_name = f"{prefix}_{suffix}"
        return self._generated_default_name

    @property
    def rom_end(self):
        return None if not self.rom else self.rom + self.size

    @property
    def vram_end(self):
        return self.vram_start + self.size

    @property
    def name(self) -> str:
        return self.given_name if self.given_name else self.default_name

    @property
    def size(self) -> int:
        if self.given_size is not None:
            return self.given_size
        return 4

    @property
    def filename(self) -> str:
        if self.given_filename is not None:
            return self.given_filename
        return self.name

    def contains_vram(self, offset):
        return offset >= self.vram_start and offset < self.vram_end

    def contains_rom(self, offset):
        return offset >= self.rom and offset < self.rom_end


def get_all_symbols():
    global all_symbols
    return all_symbols


def reset_symbols():
    global all_symbols
    global ignored_addresses
    global to_mark_as_defined
    global spim_context
    all_symbols = []
    ignored_addresses = set()
    to_mark_as_defined = set()
    spim_context = spimdisasm.common.Context()

    from .metadata import segment_manager

    segment_manager.reset()
