from __future__ import annotations

import argparse
from collections import Counter
from dataclasses import dataclass, field
from pathlib import Path

from ..formatting import h16, parse_int
from .cli import load_rom
from .constants import P9DDR, P9DR, SCI1_TDR
from .errors import UnsupportedInstruction
from .runner import H8536Emulator


DEFAULT_WATCH_PCS = (0xC08B, 0xC0DB, 0xC121, 0xBFE0, 0xBFFE, 0xC059)


@dataclass(frozen=True)
class WatchSnapshot:
    pc: int
    step: int
    regs: tuple[int, ...]
    sp: int
    stack_words: tuple[tuple[int, int], ...]
    callers: tuple[tuple[int, int | None], ...]

    def line(self) -> str:
        regs = " ".join(f"R{idx}={h16(value)}" for idx, value in enumerate(self.regs))
        stack = " ".join(f"{h16(address)}:{h16(value)}" for address, value in self.stack_words)
        if self.callers:
            callers = " ".join(
                f"{h16(return_address)}<-{h16(call_site)}" if call_site is not None else h16(return_address)
                for return_address, call_site in self.callers
            )
        else:
            callers = "-"
        return f"step={self.step} pc={h16(self.pc)} sp={h16(self.sp)} {regs} stack=[{stack}] callers=[{callers}]"


@dataclass
class ProbeReport:
    steps: int
    pc: int
    stopped_reason: str
    hot_pcs: Counter[int] = field(default_factory=Counter)
    tx_bytes: bytes = b""
    p9_bytes: list[int] = field(default_factory=list)
    p9_fast_bytes: list[int] = field(default_factory=list)
    p9_fast_events: int = 0
    p9_accesses: list[str] = field(default_factory=list)
    sci_accesses: list[str] = field(default_factory=list)
    watch_snapshots: list[WatchSnapshot] = field(default_factory=list)
    unsupported: str | None = None

    def lines(self, hot_limit: int = 12) -> list[str]:
        lines = [
            f"steps={self.steps}",
            f"pc={h16(self.pc)}",
            f"stopped={self.stopped_reason}",
            "tx_bytes=" + self.tx_bytes.hex(" ").upper(),
            "p9_bytes=" + " ".join(f"{byte:02X}" for byte in self.p9_bytes[-32:]),
            "p9_fast_bytes=" + " ".join(f"{byte:02X}" for byte in self.p9_fast_bytes[-32:]),
            f"p9_fast_events={self.p9_fast_events}",
            "hot_pcs=" + ", ".join(f"{h16(pc)}:{count}" for pc, count in self.hot_pcs.most_common(hot_limit)),
        ]
        if self.unsupported:
            lines.append(f"unsupported={self.unsupported}")
        if self.p9_accesses:
            lines.append("recent_p9:")
            lines.extend("  " + line for line in self.p9_accesses[-24:])
        if self.sci_accesses:
            lines.append("recent_sci:")
            lines.extend("  " + line for line in self.sci_accesses[-16:])
        if self.watch_snapshots:
            lines.append("recent_watch_snapshots:")
            lines.extend("  " + snapshot.line() for snapshot in self.watch_snapshots)
        return lines


def parse_watch_pc(text: str) -> int:
    try:
        value = parse_int(text)
    except ValueError:
        value = int(text, 16)
    return value & 0xFFFF


def _likely_call_site(emulator: H8536Emulator, return_address: int) -> int | None:
    rom = emulator.memory.rom
    for size in (2, 3):
        candidate = (return_address - size) & 0xFFFF
        if not rom.contains(candidate, size):
            continue
        opcode = rom.u8(candidate)
        if size == 2 and opcode == 0x0E:
            return candidate
        if size == 3 and opcode in (0x1E, 0x18):
            return candidate
    return None


def _watch_snapshot(emulator: H8536Emulator, *, stack_words: int = 6) -> WatchSnapshot:
    sp = emulator.cpu.regs[7] & 0xFFFF
    words: list[tuple[int, int]] = []
    callers: list[tuple[int, int | None]] = []
    seen_callers: set[int] = set()
    for offset in range(0, stack_words * 2, 2):
        address = (sp + offset) & 0xFFFF
        try:
            value = emulator.memory.read16(address)
        except Exception:
            continue
        words.append((address, value))
        if value in seen_callers or not emulator.memory.rom.contains(value):
            continue
        seen_callers.add(value)
        callers.append((value, _likely_call_site(emulator, value)))
    return WatchSnapshot(
        pc=emulator.cpu.pc & 0xFFFF,
        step=emulator.cpu.steps,
        regs=tuple(register & 0xFFFF for register in emulator.cpu.regs),
        sp=sp,
        stack_words=tuple(words),
        callers=tuple(callers),
    )


def run_probe(
    rom_bytes: bytes,
    *,
    max_steps: int,
    interval_steps: int,
    stop_on_tx: bool,
    p9_log_limit: int,
    frt2_ocia_steps: int = 1024,
    p9_fast_path: bool = False,
    p9_fast_input: int = 0xFF,
    watch_pcs: list[int] | tuple[int, ...] | None = None,
    watch_snapshot_limit: int = 32,
    watch_pc_limit: int = 8,
    watch_min_interval: int = 1024,
) -> ProbeReport:
    emulator = H8536Emulator(
        rom_bytes,
        interval_steps=interval_steps,
        frt2_ocia_steps=frt2_ocia_steps,
        p9_fast_path_enabled=p9_fast_path,
        p9_fast_default_input_byte=p9_fast_input,
    )
    hot_pcs: Counter[int] = Counter()
    p9_accesses: list[str] = []
    sci_accesses: list[str] = []
    snapshots: list[WatchSnapshot] = []
    watch_set = set(DEFAULT_WATCH_PCS if watch_pcs is None else watch_pcs)
    watch_counts: Counter[int] = Counter()
    watch_last_step: dict[int, int] = {}
    stopped_reason = "max_steps"
    unsupported: str | None = None
    last_access_index = 0

    for _ in range(max_steps):
        pc = emulator.cpu.pc
        hot_pcs[pc] += 1
        if pc in watch_set and watch_counts[pc] < watch_pc_limit:
            last_step = watch_last_step.get(pc)
            if last_step is None or emulator.cpu.steps - last_step >= watch_min_interval:
                snapshots.append(_watch_snapshot(emulator))
                if len(snapshots) > watch_snapshot_limit:
                    del snapshots[: len(snapshots) - watch_snapshot_limit]
                watch_counts[pc] += 1
                watch_last_step[pc] = emulator.cpu.steps
                last_access_index = len(emulator.memory.access_log)
        try:
            emulator.step()
        except UnsupportedInstruction as exc:
            stopped_reason = "unsupported_instruction"
            unsupported = str(exc)
            break

        for access in emulator.memory.access_log[last_access_index:]:
            if access.address in (P9DDR, P9DR):
                p9_accesses.append(f"{h16(pc)} {access.kind} {h16(access.address)}={access.value:02X}")
                if len(p9_accesses) > p9_log_limit:
                    del p9_accesses[: len(p9_accesses) - p9_log_limit]
            elif access.address == SCI1_TDR:
                sci_accesses.append(f"{h16(pc)} {access.kind} {h16(access.address)}={access.value:02X}")
        last_access_index = len(emulator.memory.access_log)

        if stop_on_tx and emulator.sci1.tx_bytes:
            stopped_reason = "tx"
            break
        if emulator.sci1.saw_heartbeat():
            stopped_reason = "heartbeat"
            break

    return ProbeReport(
        steps=emulator.cpu.steps,
        pc=emulator.cpu.pc,
        stopped_reason=stopped_reason,
        hot_pcs=hot_pcs,
        tx_bytes=bytes(emulator.sci1.tx_bytes),
        p9_bytes=list(emulator.memory.p9_bus.byte_candidates),
        p9_fast_bytes=list(emulator.p9_fast_path.output_bytes),
        p9_fast_events=len(emulator.p9_fast_path.events),
        p9_accesses=p9_accesses,
        sci_accesses=sci_accesses,
        watch_snapshots=snapshots,
        unsupported=unsupported,
    )


def build_arg_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Probe H8/536 emulator progress and likely hold-ups")
    parser.add_argument("--rom", type=Path, help="ROM image path; defaults to the repo ROM image")
    parser.add_argument("--max-steps", type=int, default=250_000)
    parser.add_argument("--interval-steps", type=int, default=512)
    parser.add_argument("--frt2-ocia-steps", type=int, default=1024)
    parser.add_argument("--stop-on-tx", action="store_true", help="stop when SCI1 TDR emits the first byte")
    parser.add_argument("--p9-fast-path", action="store_true", help="shortcut known P9 bit-banged transfer routines for exploration")
    parser.add_argument("--p9-fast-input", type=parse_int, default=0xFF)
    parser.add_argument("--p9-log-limit", type=int, default=80)
    parser.add_argument("--hot-limit", type=int, default=12)
    parser.add_argument(
        "--watch-pc",
        action="append",
        type=parse_watch_pc,
        default=[],
        help="additional PC to snapshot when hit, e.g. C08B, 0xC08B, or H'C08B",
    )
    parser.add_argument("--watch-snapshot-limit", type=int, default=32)
    parser.add_argument("--watch-pc-limit", type=int, default=8)
    parser.add_argument("--watch-min-interval", type=int, default=1024)
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_arg_parser().parse_args(argv)
    try:
        rom_bytes, rom_path = load_rom(args.rom)
    except FileNotFoundError as exc:
        print(str(exc))
        return 2
    print(f"rom={rom_path}")
    report = run_probe(
        rom_bytes,
        max_steps=args.max_steps,
        interval_steps=args.interval_steps,
        frt2_ocia_steps=args.frt2_ocia_steps,
        stop_on_tx=args.stop_on_tx,
        p9_log_limit=args.p9_log_limit,
        p9_fast_path=args.p9_fast_path,
        p9_fast_input=args.p9_fast_input,
        watch_pcs=tuple(dict.fromkeys((*DEFAULT_WATCH_PCS, *args.watch_pc))),
        watch_snapshot_limit=args.watch_snapshot_limit,
        watch_pc_limit=args.watch_pc_limit,
        watch_min_interval=args.watch_min_interval,
    )
    for line in report.lines(hot_limit=args.hot_limit):
        print(line)
    return 0