from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .cpu import mask, sign_bit


LOC_BFE0_TRANSFER_WRAPPER = 0xBFE0
LOC_BFFE_TRANSFER_WRAPPER = 0xBFFE
LOC_C08B_P9_WRITE_BYTE = 0xC08B
LOC_C0DB_P9_READ_BYTE = 0xC0DB
LOC_C10C_P9_MARKER = 0xC10C
LOC_C121_P9_MARKER = 0xC121
LOC_C142_P9_MARKER = 0xC142


@dataclass(frozen=True)
class P9FastPathConfig:
    """Configuration for optional ROM P9 transfer shortcuts.

    The helper assumes the CPU PC is exactly at a known routine entry. It
    models the routine as if it completed successfully and returned via RTS.
    Integration should keep this disabled unless the runner intentionally opts
    into skipping these ROM routines.
    """

    enabled: bool = False
    write_byte_pc: int = LOC_C08B_P9_WRITE_BYTE
    read_byte_pc: int = LOC_C0DB_P9_READ_BYTE
    marker_pcs: frozenset[int] = frozenset(
        {
            LOC_C10C_P9_MARKER,
            LOC_C121_P9_MARKER,
            LOC_C142_P9_MARKER,
        }
    )
    wrapper_pcs: frozenset[int] = frozenset(
        {
            LOC_BFE0_TRANSFER_WRAPPER,
            LOC_BFFE_TRANSFER_WRAPPER,
        }
    )
    default_input_byte: int = 0xFF
    account_step: bool = True
    cycles_per_hit: int = 0


@dataclass(frozen=True)
class P9FastPathEvent:
    kind: str
    pc: int
    value: int | None = None
    source: str | None = None
    queue_depth: int | None = None

    def line(self) -> str:
        value = "" if self.value is None else f" value={self.value:02X}"
        source = "" if self.source is None else f" source={self.source}"
        queue_depth = "" if self.queue_depth is None else f" queued={self.queue_depth}"
        return f"{self.kind} pc={self.pc:04X}{value}{source}{queue_depth}"


@dataclass
class P9FastPath:
    """Optional fast-path scaffold for ROM P9 bit-transfer routines."""

    config: P9FastPathConfig = field(default_factory=P9FastPathConfig)
    input_bytes: list[int] = field(default_factory=list)
    input_sources: list[str] = field(default_factory=list)
    output_bytes: list[int] = field(default_factory=list)
    events: list[P9FastPathEvent] = field(default_factory=list)

    def __post_init__(self) -> None:
        if len(self.input_sources) < len(self.input_bytes):
            self.input_sources.extend(["initial"] * (len(self.input_bytes) - len(self.input_sources)))
        elif len(self.input_sources) > len(self.input_bytes):
            del self.input_sources[len(self.input_bytes) :]

    def queue_input(self, *values: int, source: str = "queued") -> None:
        self.input_bytes.extend(value & 0xFF for value in values)
        self.input_sources.extend(source for _ in values)

    def queue_input_script(self, name: str, values: list[int] | tuple[int, ...]) -> None:
        self.queue_input(*values, source=f"script:{name}")

    def trace_lines(self, limit: int | None = None) -> list[str]:
        events = self.events if limit is None else self.events[-limit:]
        return [event.line() for event in events]

    def try_handle(self, emulator: Any) -> bool:
        if not self.config.enabled:
            return False

        pc = emulator.cpu.pc & 0xFFFF
        if pc == (self.config.write_byte_pc & 0xFFFF):
            self._handle_write_byte(emulator)
        elif pc == (self.config.read_byte_pc & 0xFFFF):
            self._handle_read_byte(emulator)
        elif pc in self.config.marker_pcs:
            self.events.append(P9FastPathEvent("marker", pc))
            self._return_from_subroutine(emulator)
        elif pc in self.config.wrapper_pcs:
            self.events.append(P9FastPathEvent("wrapper_success", pc))
            emulator.cpu.regs[0] = 1
            self._set_logic_flags(emulator.cpu, 1, 1)
            self._return_from_subroutine(emulator)
        else:
            return False

        if self.config.account_step:
            emulator.cpu.steps += 1
            emulator.cpu.cycles += self.config.cycles_per_hit
        return True

    def _handle_write_byte(self, emulator: Any) -> None:
        pc = emulator.cpu.pc & 0xFFFF
        value = emulator.cpu.regs[0] & 0xFF
        self.output_bytes.append(value)
        self.events.append(P9FastPathEvent("write_byte", pc, value))

        emulator.cpu.regs[0] = 1
        self._set_logic_flags(emulator.cpu, 1, 1)
        self._return_from_subroutine(emulator)

    def _handle_read_byte(self, emulator: Any) -> None:
        pc = emulator.cpu.pc & 0xFFFF
        if self.input_bytes:
            value = self.input_bytes.pop(0)
            source = self.input_sources.pop(0) if self.input_sources else "queued"
        else:
            value = self.config.default_input_byte
            source = "default_input_byte"
        value &= 0xFF
        self.events.append(P9FastPathEvent("read_byte", pc, value, source, len(self.input_bytes)))

        # The ROM-side read routine yields a byte in R5. Model that as a byte
        # register write so the existing high byte is not accidentally clobbered.
        emulator.cpu.regs[5] = (emulator.cpu.regs[5] & 0xFF00) | value
        self._set_logic_flags(emulator.cpu, value, 1)
        self._return_from_subroutine(emulator)

    def _return_from_subroutine(self, emulator: Any) -> None:
        sp = emulator.cpu.regs[7] & 0xFFFF
        emulator.cpu.pc = emulator.memory.read16(sp) & 0xFFFF
        emulator.cpu.regs[7] = (sp + 2) & 0xFFFF

    def _set_logic_flags(self, cpu: Any, value: int, size: int) -> None:
        value &= mask(size)
        cpu.z = value == 0
        cpu.n = bool(value & sign_bit(size))
        cpu.v = False