"""Native risc-v-gsp-rm-debugger execution backend (M1): deterministic RV64IMAC in Python.

``ge.launch(kernel=..., backend="native")`` builds a :class:`SimTarget` around
:class:`sim.machine.Machine` and exposes a :class:`SimClient` whose method
surface matches :class:`greeneye.rsp.RspClient`, so ``greeneye.session.Session``
runs unchanged against it.  ``trace=`` enables the M1.5 exec/access trace
(:mod:`sim.trace`), written to the session as ``blobs/trace.jsonl``.
``device_map=`` and ``unknown_mmio="lenient"`` (H2.1/H2.2 start) serve
unmodelled MMIO accesses from a harvested first-touch map instead of faulting;
lenient accesses are logged to ``target.mmio_log`` and counted in
``target.mmio_unknown``.

Virtual time is instruction count: ``continue_(timeout)`` executes at most
``timeout * SIM_IPS`` instructions (SIM_IPS = 1 MIPS), deterministically.  The
wall-clock budgets ``Session.run`` derives from its deadline therefore bound
the run; the stop/watch/hook sequence within a fixed instruction budget is
reproducible on every host.
"""
from __future__ import annotations

import os
import sys

from .elf import load_elf
from .errors import GreenEyeError, RspProtocolError, RspTimeout
from .rsp import Register, StopReply
from .session import Session
from .symbols import SymbolDB

_TOOL_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _TOOL_ROOT not in sys.path:  # `sim/` lives beside the greeneye package
    sys.path.insert(0, _TOOL_ROOT)

from sim.machine import (  # noqa: E402
    R_BUDGET,
    R_FAULT,
    R_WATCH,
    Machine,
)
from sim.trace import make_trace  # noqa: E402
#: Virtual instructions per second used to convert wall timeouts to budgets.
SIM_IPS = 1_000_000
#: Safety valve for ``continue_(timeout=None)``: never run unbounded.
MAX_UNBOUNDED_BUDGET = 20_000_000

#: Canonical register names, as exposed by QEMU's target.xml for riscv:rv64.
ABI_NAMES = (
    ["zero", "ra", "sp", "gp", "tp", "t0", "t1", "t2", "s0", "s1"]
    + [f"a{i}" for i in range(8)]
    + [f"s{i}" for i in range(2, 12)]
    + [f"t{i}" for i in range(3, 7)]
)
#: xN -> ABI-name aliases (x0 == zero, x1 == ra, ...).
_X_TO_ABI = {f"x{i}": ABI_NAMES[i] for i in range(32)}
_X_TO_ABI["fp"] = "s0"
_X_TO_ABI["return"] = "ra"
_ABI_TO_INDEX = {name: i for i, name in enumerate(ABI_NAMES)}


class SimClient:
    """Drop-in replacement for :class:`RspClient` over a native Machine."""

    def __init__(self, machine: Machine, *, sim_ips: int = SIM_IPS, trace=None):
        self.machine = machine
        self.sim_ips = max(1, int(sim_ips))
        self.timeout = 10.0
        self.packet_size = 0x1000
        self.trace = trace  # sim.trace.Trace | None
        self._recorder = None
        self.registers = [
            Register(number=i, name=name, bitsize=64, offset=i * 8)
            for i, name in enumerate(ABI_NAMES)
        ] + [Register(number=32, name="pc", bitsize=64, offset=32 * 8)]
        self._by_name = {r.name: r for r in self.registers}
        self._by_name.update(_X_TO_ABI)
        self._by_number = {r.number: r for r in self.registers}
        self.features = "PacketSize=1000;swbreak+;hwbreak+;greeneye-native+"
        self.connected = True
        self._bps: set[int] = set()
        self._hw_bps: set[int] = set()
        self._watches: dict[tuple[int, int, str], None] = {}

    # -- lifecycle ---------------------------------------------------------

    def connect(self, retries: int = 1, retry_delay: float = 0.1) -> None:
        self.connected = True

    def sync(self) -> None:
        pass  # no server to hand-shake with

    def close(self) -> None:
        self.connected = False
        if self.trace is not None:
            self.trace.flush(self._recorder)

    def __enter__(self) -> "SimClient":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    # -- registers ---------------------------------------------------------

    def read_registers(self) -> dict[str, int]:
        x = self.machine.x
        out = {name: x[i] for i, name in enumerate(ABI_NAMES)}
        out["fp"] = out["s0"]  # QEMU's target.xml names x8 "fp"
        out["pc"] = self.machine.pc
        return out

    def _index(self, name_or_number: str | int) -> int:
        if isinstance(name_or_number, int):
            if 0 <= name_or_number < 32:
                return name_or_number
            if name_or_number == 32:
                return 32
            raise RspProtocolError(f"unknown register number {name_or_number}")
        key = _X_TO_ABI.get(name_or_number, name_or_number)
        if key == "pc":
            return 32
        index = _ABI_TO_INDEX.get(key)
        if index is None:
            raise RspProtocolError(f"unknown register {name_or_number!r}")
        return index

    def read_reg(self, name_or_number: str | int) -> int:
        index = self._index(name_or_number)
        if index == 32:
            return self.machine.pc
        return self.machine.reg(index)

    def write_reg(self, name_or_number: str | int, value: int) -> None:
        index = self._index(name_or_number)
        if index == 32:
            self.machine.pc = value & ((1 << 64) - 1)
        else:
            self.machine.set_reg(index, value)

    @property
    def pc(self) -> int:
        return self.machine.pc

    # -- memory ------------------------------------------------------------

    def read_mem(self, addr: int, size: int, chunk: int = 0) -> bytes:
        try:
            return self.machine.read_mem(addr, size)
        except Exception as exc:
            raise RspProtocolError(str(exc)) from exc

    def write_mem(self, addr: int, data: bytes, chunk: int = 0) -> None:
        try:
            self.machine.write_mem(addr, data)
        except Exception as exc:
            raise RspProtocolError(str(exc)) from exc

    def read_cstring(self, addr: int, max_len: int = 256) -> str:
        raw = self.read_mem(addr, max_len)
        return raw.split(b"\x00", 1)[0].decode("utf-8", "replace")

    # -- breakpoints / watchpoints ----------------------------------------

    def set_sw_break(self, addr: int, kind: int = 4) -> None:
        self._bps.add(addr)

    def clear_sw_break(self, addr: int, kind: int = 4) -> None:
        self._bps.discard(addr)

    def set_hw_break(self, addr: int, kind: int = 4) -> None:
        self._hw_bps.add(addr)

    def clear_hw_break(self, addr: int, kind: int = 4) -> None:
        self._hw_bps.discard(addr)

    _WATCH_KINDS = {"w": 2, "r": 3, "a": 4}
    _WATCH_REVERSE = {"w": "watch", "r": "rwatch", "a": "awatch"}

    def set_watch(self, addr: int, size: int = 8, kind: str = "w") -> None:
        if kind not in self._WATCH_KINDS:
            raise RspProtocolError(f"unknown watchpoint kind {kind!r}")
        self._watches[(addr, size, kind)] = None

    def clear_watch(self, addr: int, size: int = 8, kind: str = "w") -> None:
        self._watches.pop((addr, size, kind), None)

    def _watcher_list(self):
        out = []
        for (addr, size, kind) in self._watches:
            mask = (2 if kind in ("r", "a") else 0) | (1 if kind in ("w", "a") else 0)
            out.append((addr, addr + size, kind, mask))
        return tuple(out)

    # -- run control -------------------------------------------------------

    def _budget(self, timeout: float | None) -> int | None:
        if timeout is None:
            return None
        return max(1, int(timeout * self.sim_ips))

    def _to_stop(self, result) -> StopReply:
        reason, pc, a, b = result
        if reason == R_WATCH:
            kind = self._WATCH_REVERSE.get(b, "watch")
            return StopReply(
                raw=f"T05watch:{a:x};",
                signal=5,
                watch_addr=a,
                watch_kind=kind,
            )
        if reason == R_FAULT:
            signal = 4 if a == 2 else 11
            return StopReply(raw=f"T0{signal:x}", signal=signal)
        # bp, ebreak, ecall and step all report SIGTRAP, like the gdbstub.
        return StopReply(raw="T05", signal=5)

    def continue_(self, timeout: float | None = None) -> StopReply:
        budget = self._budget(timeout)
        if budget is None:
            budget = MAX_UNBOUNDED_BUDGET
        result = self.machine.run(
            budget,
            breakpoints=frozenset(self._bps | self._hw_bps),
            watchers=self._watcher_list(),
            trace=self.trace,
        )
        if result[0] == R_BUDGET:
            raise RspTimeout(f"native budget exhausted after {budget} instructions")
        return self._to_stop(result)

    def step(self, timeout: float | None = None) -> StopReply:
        result = self.machine.step(watchers=self._watcher_list(), trace=self.trace)
        return self._to_stop(result)

    def interrupt(self, timeout: float | None = None) -> StopReply:
        return StopReply(raw="T02", signal=2)

    # -- monitor -----------------------------------------------------------

    def monitor(self, cmd: str) -> str:
        return ""


class SimTarget:
    """QemuTarget-shaped handle for a native Machine (no process, no socket)."""

    def __init__(
        self,
        machine: Machine,
        *,
        memory_mb: int = 256,
        ram_base: int = 0x80000000,
        sim_ips: int = SIM_IPS,
        keep_alive: bool = False,
        trace=None,
    ):
        self.machine = machine
        self.memory_mb = memory_mb
        self.ram_base = ram_base
        self.keep_alive = keep_alive
        self.client = SimClient(machine, sim_ips=sim_ips, trace=trace)
        self.kernel: str | None = None
        self.workdir: str | None = None
        self._alive = True

    # QemuTarget lifecycle surface ----------------------------------------

    def start(self, paused: bool = True, timeout: float = 0.0) -> "SimTarget":
        return self

    def wait_for_port(self, timeout: float = 0.0) -> None:
        return None

    def alive(self) -> bool:
        return self._alive

    def kill(self, timeout: float = 0.0) -> None:
        self._alive = False
        self.client.close()

    def read_uart(self) -> str:
        return self.machine.uart.text()

    @property
    def mmio_log(self) -> list[dict]:
        """Lenient-mode log (empty unless ``unknown_mmio="lenient"``)."""
        return self.machine.mmio_log

    @property
    def mmio_unknown(self) -> int:
        """Count of unmodelled MMIO accesses served leniently."""
        return self.machine.mmio_unknown

    def wait_for_uart(self, marker: str, timeout: float = 5.0, interval: float = 0.05) -> bool:
        # Execution is synchronous: by the time the caller looks, the marker
        # either exists or the budget was too small.  No polling needed.
        return marker in self.read_uart()

    def read_qemu_log(self) -> str:
        return ""

    @property
    def command_line(self) -> list[str]:
        return [f"native-sim:{self.kernel or '<none>'}"]

    def shell_command(self) -> str:
        return " ".join(self.command_line)

    def __repr__(self) -> str:  # pragma: no cover - cosmetic
        return f"SimTarget(pc={self.machine.pc:#x}, alive={self._alive})"


def launch_native(
    *,
    kernel: str | None = None,
    memory_mb: int = 256,
    ram_base: int = 0x80000000,
    record=None,
    symbols: SymbolDB | None = None,
    name: str = "session",
    load_symbols: bool = True,
    keep_alive: bool = False,
    sim_ips: int = SIM_IPS,
    trace=None,
    device_map=None,
    unknown_mmio: str = "fault",
    learn_mmio: bool = False,
    bias: int = 0,
    cycle_model=None,
) -> Session:
    """Boot the native simulator and return a ready Session.

    ``trace`` enables the R4 execution/access trace: ``True`` for defaults, a
    dict of ``{"exec": bool, "start": int, "end": int, "max_events": int}``,
    or a pre-built :class:`sim.trace.Trace`.  The trace is written to the
    session as ``blobs/trace.jsonl`` when the Session closes; tracing implies
    recording, so a session directory is created when none was given.

    ``device_map`` is a :class:`greeneye.device_map.DeviceMap` (or a path to a
    saved map) describing modelled MMIO; ``unknown_mmio="lenient"`` serves
    unmodelled accesses from it (or zero) instead of faulting, logging each in
    ``machine.mmio_log`` and counting ``machine.mmio_unknown``.
    ``learn_mmio=True`` writes the observed values back into the map.
    ``bias`` relocates PT_LOAD segments and the entry for images linked for
    GPU physical addresses (e.g. ``0x80000000`` for the real GSP fwimage).

    ``cycle_model`` selects the modeled instruction-class cycle accounting
    (:mod:`sim.cycles`): ``None`` (default) uses the default table, ``False``
    disables it (``machine.cycles`` stays 0), or pass a ``CycleModel``.  The
    model never changes execution.
    """
    try:
        tracer = make_trace(trace)
    except ValueError as exc:
        raise GreenEyeError(f"invalid trace spec: {exc}") from exc
    if tracer is not None and not record:
        record = True  # the trace artifact lives in the session
    if isinstance(device_map, (str, os.PathLike)):
        from .device_map import DeviceMap

        device_map = DeviceMap.load(os.fspath(device_map))
    machine = Machine(
        memory_mb=memory_mb,
        ram_base=ram_base,
        device_map=device_map,
        unknown_mmio=unknown_mmio,
        learn_mmio=learn_mmio,
        bias=bias,
        cycle_model=cycle_model,
    )
    target = SimTarget(
        machine,
        memory_mb=memory_mb,
        ram_base=ram_base,
        sim_ips=sim_ips,
        keep_alive=keep_alive,
        trace=tracer,
    )
    target.kernel = kernel
    symdb = symbols or SymbolDB()
    meta = {"backend": "native", "sim_ips": sim_ips, "firmware": kernel}
    if unknown_mmio != "fault" or device_map is not None:
        meta["unknown_mmio"] = unknown_mmio
    session = Session(
        client=target.client,
        target=target,
        symbols=symdb,
        record=record,
        name=name,
        meta=meta,
    )
    if tracer is not None:
        target.client._recorder = session.recorder
        session.recorder.meta["trace"] = {**tracer.config()}
    if kernel:
        image = load_elf(kernel)
        if not image.is_riscv:
            raise GreenEyeError(f"{kernel}: not a RISC-V ELF image")
        machine.load_segments(image.flat_load_plan(bias=bias), image.relocated_entry(bias=bias))
        session.load_image(image, bias=bias, load_symbols=load_symbols)
    return session