"""Native execution/access trace (roadmap R4): the in-process equivalent of
``qemu/greeneye_trace.c``.

The QEMU TCG plugin records every translated instruction and every memory
access as JSONL.  This module does the same from inside
:class:`sim.machine.Machine`, with the same knobs and line shape:

* ``exec=on|off`` - emit one ``exec`` line per executed instruction;
* ``start=``/``end=`` - only instrument PCs in the inclusive range (matching
  the plugin, which filters at translation time, so memory accesses are also
  filtered by the *instruction* PC);
* a bounded ring budget - once full, further lines are counted as ``dropped``
  instead of growing without limit (deterministically).

Lines are JSON objects in the TCG plugin's schema (``t``/``pc``/``addr``/
``size``, hex strings), plus a ``vt`` virtual-time field (the machine's
instruction count).  Nothing here reads the host clock, so the payload is
byte-identical for identical execution.
"""
from __future__ import annotations

#: Default cap on buffered trace lines (~50 bytes each in memory).
DEFAULT_MAX_EVENTS = 1_000_000


class Trace:
    """Collects compact exec/access lines for one native session."""

    def __init__(self, *, exec_enabled: bool = True, start: int | None = None,
                 end: int | None = None, max_events: int | None = None):
        self.exec_enabled = bool(exec_enabled)
        self.start = start
        self.end = end
        self.max_events = DEFAULT_MAX_EVENTS if max_events is None else max(0, int(max_events))
        self.execs = 0
        self.reads = 0
        self.writes = 0
        self.events = 0
        self.dropped = 0
        self._lines: list[str] = []
        self._flushed = False

    # -- configuration -----------------------------------------------------

    def in_range(self, pc: int) -> bool:
        if self.start is not None and pc < self.start:
            return False
        if self.end is not None and pc > self.end:
            return False
        return True

    def config(self) -> dict:
        return {
            "exec": self.exec_enabled,
            "start": self.start,
            "end": self.end,
            "max_events": self.max_events,
        }

    def counts(self) -> dict:
        return {
            "exec": self.execs,
            "read": self.reads,
            "write": self.writes,
            "events": self.events,
            "dropped": self.dropped,
        }

    # -- collection (called from Machine.run's hot loop) --------------------

    def begin(self, pc: int, vt: int) -> bool:
        """Called once per executed instruction; returns access-capture."""
        if not self.in_range(pc):
            return False
        if self.exec_enabled:
            self._append(f'{{"t":"exec","vcpu":0,"pc":"0x{pc:x}","vt":{vt}}}')
            self.execs += 1
        return True

    def mem(self, pc: int, addr: int, size: int, is_write: bool, vt: int) -> None:
        kind = "write" if is_write else "read"
        if is_write:
            self.writes += 1
        else:
            self.reads += 1
        self._append(
            f'{{"t":"{kind}","vcpu":0,"addr":"0x{addr:x}","size":{size},'
            f'"pc":"0x{pc:x}","vt":{vt}}}'
        )

    def _append(self, line: str) -> None:
        self.events += 1
        if self.max_events and len(self._lines) >= self.max_events:
            self.dropped += 1
            return
        self._lines.append(line)

    # -- output ------------------------------------------------------------

    def payload(self) -> bytes:
        """The complete JSONL document, including the plugin's summary line."""
        body = "\n".join(self._lines)
        if body:
            body += "\n"
        summary = (
            f'{{"t":"summary","insns":{self.execs},"mem":{self.reads + self.writes},'
            f'"events":{self.events},"dropped":{self.dropped},'
            f'"exec":{str(self.exec_enabled).lower()},'
            f'"start":{_json_addr(self.start)},"end":{_json_addr(self.end)}}}'
        )
        return (body + summary + "\n").encode("utf-8")

    def flush(self, recorder) -> str | None:
        """Write the trace to ``recorder`` as ``blobs/trace.jsonl`` (once)."""
        if self._flushed:
            return None
        self._flushed = True
        if recorder is None or getattr(recorder, "dir", None) is None:
            return None
        rel = recorder.blob("trace.jsonl", self.payload())
        recorder.meta["trace_blob"] = rel
        recorder.meta["trace"] = {**self.config(), **self.counts()}
        return rel


def make_trace(spec):
    """Normalise a trace spec (None/False, True, dict or Trace) to a Trace."""
    if spec is None or spec is False:
        return None
    if isinstance(spec, Trace):
        return spec
    if spec is True:
        return Trace()
    if isinstance(spec, dict):
        return Trace(
            exec_enabled=spec.get("exec", spec.get("exec_enabled", True)),
            start=spec.get("start"),
            end=spec.get("end"),
            max_events=spec.get("max_events", spec.get("max")),
        )
    raise ValueError(f"trace spec must be None, bool, dict or Trace, not {type(spec).__name__}")


def _json_addr(value: int | None) -> str:
    return "null" if value is None else f"\"0x{value:x}\""