from __future__ import annotations

from collections.abc import Iterable

from .decoder import H8536Decoder
from .formatting import label_for
from .model import Instruction
from .tables import FLOW_STOP


def trace(decoder: H8536Decoder, starts: Iterable[int], start_limit: int, end_limit: int) -> dict[int, Instruction]:
    decoded: dict[int, Instruction] = {}
    queue = [(addr, decoder.br) for addr in starts if start_limit <= addr < end_limit]
    seen_starts: set[tuple[int, int | None]] = set()

    while queue:
        pc, br = queue.pop(0)
        if (pc, br) in seen_starts:
            continue
        seen_starts.add((pc, br))

        while start_limit <= pc < end_limit:
            decoder.br = br
            if pc in decoded:
                break
            ins = decoder.decode(pc)
            decoded[pc] = ins
            next_br = ins.br_value if ins.writes_br else br
            for target in ins.targets:
                if start_limit <= target < end_limit and (target, next_br) not in seen_starts:
                    queue.append((target, next_br))
            if ins.kind in FLOW_STOP or not ins.fallthrough:
                break
            pc = (pc + max(ins.size, 1)) & 0xFFFF
            br = next_br

    return decoded


def linear_sweep(decoder: H8536Decoder, start: int, end: int) -> dict[int, Instruction]:
    decoded: dict[int, Instruction] = {}
    pc = start
    br = decoder.br
    while pc < end:
        decoder.br = br
        ins = decoder.decode(pc)
        decoded[pc] = ins
        if ins.writes_br:
            br = ins.br_value
        pc += max(ins.size, 1)
    return decoded


def collect_labels(instructions: Iterable[Instruction], vectors: dict[int, tuple[str, int]]) -> dict[int, str]:
    labels: dict[int, str] = {}
    for _vector_addr, (name, target) in vectors.items():
        labels.setdefault(target, f"vec_{name}_{target:04X}")
    for ins in instructions:
        for target in ins.targets:
            labels.setdefault(target, label_for(target))
    return labels


def collect_function_entries(
    instructions: Iterable[Instruction],
    vectors: dict[int, tuple[str, int]],
) -> set[int]:
    entries = {target for _name, target in vectors.values()}
    for ins in instructions:
        if ins.kind == "call":
            entries.update(ins.targets)
    return entries


def assign_functions(instructions: dict[int, Instruction], entries: set[int]) -> dict[int, int]:
    owners: dict[int, int] = {}
    current: int | None = None
    for address in sorted(instructions):
        if address in entries:
            current = address
        if current is not None:
            owners[address] = current
        if instructions[address].kind in {"return", "rte", "sleep"}:
            current = None
    return owners


def build_functions(
    instructions: dict[int, Instruction],
    vectors: dict[int, tuple[str, int]],
    labels: dict[int, str],
) -> list[dict[str, object]]:
    entries = collect_function_entries(instructions.values(), vectors)
    owners = assign_functions(instructions, entries)
    vector_sources: dict[int, list[str]] = {}
    for _vector_addr, (name, target) in vectors.items():
        vector_sources.setdefault(target, []).append(name)

    functions: dict[int, dict[str, object]] = {}
    for address, owner in owners.items():
        ins = instructions[address]
        function = functions.setdefault(
            owner,
            {
                "start": owner,
                "label": labels.get(owner, label_for(owner)),
                "sources": vector_sources.get(owner, []),
                "instruction_count": 0,
                "end": owner,
                "calls": [],
                "unresolved_calls": 0,
            },
        )
        function["instruction_count"] = int(function["instruction_count"]) + 1
        function["end"] = max(int(function["end"]), ins.address + max(ins.size, 1) - 1)
        if ins.kind == "call":
            if ins.targets:
                calls = function["calls"]
                assert isinstance(calls, list)
                for target in ins.targets:
                    if target not in calls:
                        calls.append(target)
            else:
                function["unresolved_calls"] = int(function["unresolved_calls"]) + 1

    return [functions[start] for start in sorted(functions)]


def build_call_graph(
    instructions: dict[int, Instruction],
    vectors: dict[int, tuple[str, int]],
    labels: dict[int, str],
) -> dict[str, object]:
    entries = collect_function_entries(instructions.values(), vectors)
    owners = assign_functions(instructions, entries)
    nodes = build_functions(instructions, vectors, labels)
    edges: list[dict[str, object]] = []
    seen: set[tuple[int, int]] = set()

    for ins in instructions.values():
        if ins.kind != "call" or not ins.targets:
            continue
        source = owners.get(ins.address)
        if source is None:
            continue
        for target in ins.targets:
            key = (source, target)
            if key in seen:
                continue
            seen.add(key)
            edges.append(
                {
                    "from": source,
                    "from_label": labels.get(source, label_for(source)),
                    "to": target,
                    "to_label": labels.get(target, label_for(target)),
                    "call_site": ins.address,
                },
            )

    return {"nodes": nodes, "edges": sorted(edges, key=lambda edge: (edge["from"], edge["to"]))}