"""Versioned first-touch device map (roadmap H2.1 -> H2.2 start).

A device map is JSON data, not code: it records which MMIO addresses a
firmware image touched first, at which width and direction, with observed
values, hit counts and provenance.  ``harvest_*`` builds one from the QEMU
logs a ``greeneye probe`` session records; the native simulator
(:class:`sim.machine.Machine`) consults it so unmodelled accesses return a
mapped behaviour instead of faulting (``unknown_mmio="lenient"``).

Schema ``greeneye-device-map/1``::

    {
      "schema": "greeneye-device-map/1",
      "meta": {"bias": 2147483648, "firmware": "...", "firmware_sha256": "..."},
      "entries": [
        {"address": 2305843009213696512, "width": 4, "kind": "r",
         "values": [], "hits": 1, "first_epc": 2214629396,
         "first_epc_fw": 67145716, "source": "probe", "notes": "fault_load"},
        ...
      ]
    }

Entries are emitted in a deterministic order (address, width, kind) and the
file contains no timestamps or host-specific paths in semantic fields, so two
harvests of the same log are byte-identical.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
from dataclasses import dataclass

from .errors import GreenEyeError

SCHEMA = "greeneye-device-map/1"
KINDS = ("r", "w")
SOURCES = ("probe", "manual", "trace")
_RAM_BASE = 0x80000000
#: QEMU probe first fault for the TU104 570.144 image (docs/06-ACCURACY.md §3).
QEMU_REFERENCE = {
    "firmware_vaddr": 0x4009014,
    "access_addr": 0x2000000000000A00,
    "size": 4,
    "rw": "read",
    "cause": 5,
    "desc": "load access fault",
}

_BEHAVIOUR_ZERO = "zero"
_BEHAVIOUR_VALUE = "value:"


def _fw_pc(pc: int, bias: int) -> int:
    """Translate a relocated PC back to the firmware's own vaddr.

    A PC of 0 is a trap-vector fetch (not relocated code) and stays 0; PCs
    below the relocation bias are not attributable and stay untranslated.
    """
    if not bias or pc == 0 or pc < bias:
        return pc
    return pc - bias


@dataclass(frozen=True)
class DeviceEntry:
    """One observed (address, width, direction) MMIO touch."""

    address: int
    width: int
    kind: str  # "r" | "w"
    values: tuple[int, ...] = ()
    hits: int = 1
    first_epc: int = 0  # PC as observed in the run (relocated)
    first_epc_fw: int = 0  # PC - bias: the firmware's own virtual address
    source: str = "manual"  # probe | manual | trace
    notes: str = ""

    def key(self) -> tuple[int, int, str]:
        return (self.address, self.width, self.kind)

    def to_dict(self) -> dict:
        return {
            "address": self.address,
            "width": self.width,
            "kind": self.kind,
            "values": list(self.values),
            "hits": self.hits,
            "first_epc": self.first_epc,
            "first_epc_fw": self.first_epc_fw,
            "source": self.source,
            "notes": self.notes,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "DeviceEntry":
        return cls(
            address=int(data["address"]),
            width=int(data["width"]),
            kind=str(data["kind"]),
            values=tuple(int(v) for v in data.get("values", ())),
            hits=int(data.get("hits", 1)),
            first_epc=int(data.get("first_epc", 0)),
            first_epc_fw=int(data.get("first_epc_fw", 0)),
            source=str(data.get("source", "manual")),
            notes=str(data.get("notes", "")),
        )


class DeviceMap:
    """A versioned collection of :class:`DeviceEntry` records."""

    def __init__(self, *, meta: dict | None = None):
        self.meta: dict = dict(meta or {})
        self._entries: dict[tuple[int, int, str], DeviceEntry] = {}

    # -- construction ------------------------------------------------------

    def record(
        self,
        addr: int,
        size: int,
        kind: str,
        value: int | None = None,
        *,
        pc: int = 0,
        bias: int = 0,
        source: str = "manual",
        notes: str = "",
    ) -> DeviceEntry:
        """Add or merge one observation; returns the resulting entry."""
        if kind not in KINDS:
            raise GreenEyeError(f"device map: access kind must be r or w, got {kind!r}")
        if size <= 0:
            raise GreenEyeError(f"device map: width must be positive, got {size}")
        key = (addr, size, kind)
        entry = self._entries.get(key)
        if entry is None:
            entry = DeviceEntry(
                address=addr,
                width=size,
                kind=kind,
                values=(value,) if value is not None else (),
                hits=1,
                first_epc=pc,
                first_epc_fw=_fw_pc(pc, bias),
                source=source,
                notes=notes,
            )
        else:
            values = entry.values
            if value is not None and value not in values:
                values = values + (value,)
            entry = DeviceEntry(
                address=entry.address,
                width=entry.width,
                kind=entry.kind,
                values=values,
                hits=entry.hits + 1,
                first_epc=entry.first_epc or pc,
                first_epc_fw=entry.first_epc_fw or _fw_pc(pc, bias),
                source=entry.source,
                notes=entry.notes or notes,
            )
        self._entries[key] = entry
        return entry

    def add(self, entry: DeviceEntry) -> DeviceEntry:
        """Merge a whole entry (used by :meth:`merge`)."""
        key = entry.key()
        current = self._entries.get(key)
        if current is None:
            self._entries[key] = entry
            return entry
        values = current.values
        for value in entry.values:
            if value not in values:
                values = values + (value,)
        merged = DeviceEntry(
            address=current.address,
            width=current.width,
            kind=current.kind,
            values=values,
            hits=current.hits + entry.hits,
            first_epc=current.first_epc or entry.first_epc,
            first_epc_fw=current.first_epc_fw or entry.first_epc_fw,
            source=current.source,
            notes=current.notes or entry.notes,
        )
        self._entries[key] = merged
        return merged

    def merge(self, other: "DeviceMap") -> "DeviceMap":
        """Merge ``other`` into this map (deterministic given call order)."""
        for key in sorted(other._entries):
            self.add(other._entries[key])
        for key, value in other.meta.items():
            self.meta.setdefault(key, value)
        return self

    # -- queries -----------------------------------------------------------

    def entries(self) -> list[DeviceEntry]:
        """All entries, sorted by ``(address, width, kind)``."""
        return [self._entries[key] for key in sorted(self._entries)]

    def __len__(self) -> int:
        return len(self._entries)

    def _find(self, addr: int, size: int, kind: str) -> DeviceEntry | None:
        exact = self._entries.get((addr, size, kind))
        if exact is not None:
            return exact
        for key in sorted(self._entries):  # width fallback: smallest width first
            k_addr, _k_size, k_kind = key
            if k_addr == addr and k_kind == kind:
                return self._entries[key]
        return None

    def lookup(self, addr: int, size: int, write: bool = False) -> str:
        """Behaviour spec for an access: ``zero``, ``value:0x...`` or ``unimplemented``.

        Exact (address, width, direction) wins; a different width at the same
        address is still evidence for the register.  A write to an address
        known only for reads is absorbed (``zero``); a read of a write-only
        register reads back zero.  An address with no entry at all is
        ``unimplemented``.
        """
        kind = "w" if write else "r"
        entry = self._find(addr, size, kind)
        if entry is None:
            if self._find(addr, size, "w" if not write else "r") is not None:
                return _BEHAVIOUR_ZERO  # known address, opposite direction
            return "unimplemented"
        if not write and entry.values:
            return f"{_BEHAVIOUR_VALUE}0x{entry.values[-1]:x}"
        return _BEHAVIOUR_ZERO

    # -- persistence -------------------------------------------------------

    def to_dict(self) -> dict:
        return {
            "schema": SCHEMA,
            "meta": dict(self.meta),
            "entries": [entry.to_dict() for entry in self.entries()],
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), indent=1, sort_keys=True) + "\n"

    @classmethod
    def from_dict(cls, data: dict) -> "DeviceMap":
        schema = data.get("schema")
        if schema != SCHEMA:
            raise GreenEyeError(f"device map: unsupported schema {schema!r} (want {SCHEMA!r})")
        dm = cls(meta=data.get("meta") or {})
        for raw in data.get("entries", ()):
            dm.add(DeviceEntry.from_dict(raw))
        return dm

    @classmethod
    def from_json(cls, text: str) -> "DeviceMap":
        try:
            data = json.loads(text)
        except ValueError as exc:
            raise GreenEyeError(f"device map: invalid JSON ({exc})") from exc
        if not isinstance(data, dict):
            raise GreenEyeError("device map: top level must be an object")
        return cls.from_dict(data)

    def save(self, path: str) -> str:
        with open(path, "w", encoding="utf-8") as fh:
            fh.write(self.to_json())
        return path

    @classmethod
    def load(cls, path: str) -> "DeviceMap":
        try:
            with open(path, "r", encoding="utf-8") as fh:
                text = fh.read()
        except OSError as exc:
            raise GreenEyeError(f"device map: cannot read {path}: {exc}") from exc
        return cls.from_json(text)


def firmware_sha256(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


# -- harvest ---------------------------------------------------------------


def harvest_from_text(log_text: str, bias: int = 0, source: str = "probe") -> DeviceMap:
    """Build a map from an arbitrary QEMU ``guest_errors,unimp,int`` log.

    Only records that pair an invalid access with its exception contribute an
    entry, using the faulting PC (translated by ``bias``) and the access
    address/width/direction.  Fetch-loop faults (no access) are skipped.
    """
    from .analyze import parse_qemu_faults

    dm = DeviceMap(meta={"bias": bias, "source": source})
    for fault in parse_qemu_faults(log_text):
        access = fault.get("access")
        if not access:
            continue
        kind = "r" if access.get("rw") == "read" else "w"
        desc = fault.get("desc") or ""
        reason = access.get("reason") or ""
        dm.record(
            int(access["access_addr"]),
            int(access["size"]),
            kind,
            pc=int(fault.get("epc", 0)),
            bias=bias,
            source=source,
            notes=f"{desc} reason={reason}".strip(),
        )
    return dm


def _fault_to_map(fault: dict, bias: int, source: str = "probe") -> DeviceMap:
    dm = DeviceMap(meta={"bias": bias, "source": source})
    access = fault.get("access")
    if access:
        kind = "r" if access.get("rw") == "read" else "w"
        dm.record(
            int(access["access_addr"]),
            int(access["size"]),
            kind,
            pc=int(fault.get("epc", 0)),
            bias=bias,
            source=source,
            notes=str(fault.get("desc") or ""),
        )
    return dm


def harvest_probe_session(session_dir: str, bias: int | None = None) -> DeviceMap:
    """Harvest a ``greeneye probe`` session directory into a device map.

    Reads ``session.json`` for the relocation bias and firmware identity, then
    the ``qemu_log_blob`` (falling back to ``blobs/qemu.log``, ``qemu.log`` or
    the stored ``qemu_first_fault`` record when no blob was kept).  ``bias``
    overrides the session's recorded relocation bias when given.
    """
    session_file = os.path.join(session_dir, "session.json")
    try:
        with open(session_file, "r", encoding="utf-8") as fh:
            session = json.load(fh)
    except FileNotFoundError as exc:
        raise GreenEyeError(f"not a greeneye session directory: {session_dir}") from exc
    except ValueError as exc:
        raise GreenEyeError(f"{session_file}: invalid JSON ({exc})") from exc

    if bias is None:
        bias = int(session.get("bias", 0) or 0)
    log_text: str | None = None
    blob = session.get("qemu_log_blob")
    candidates = [blob] if blob else []
    candidates += ["blobs/qemu.log", "qemu.log"]
    log_path: str | None = None
    for candidate in candidates:
        if candidate and os.path.exists(os.path.join(session_dir, candidate)):
            log_path = os.path.join(session_dir, candidate)
            break
    if log_path is not None:
        with open(log_path, "r", encoding="utf-8", errors="replace") as fh:
            log_text = fh.read()

    first = session.get("qemu_first_fault")
    if log_text is not None:
        dm = harvest_from_text(log_text, bias=bias)
    elif isinstance(first, dict):
        dm = _fault_to_map(first, bias)
    else:
        dm = DeviceMap(meta={"bias": bias})

    meta = dm.meta
    meta["session"] = os.path.basename(os.path.abspath(session_dir))
    if session.get("firmware"):
        meta["firmware"] = session["firmware"]
    if session.get("entry") is not None:
        meta["entry"] = int(session["entry"])
    if session.get("qemu_fault_count") is not None:
        meta["qemu_fault_count"] = int(session["qemu_fault_count"])
    fw = session.get("firmware")
    if fw and os.path.isfile(fw):
        meta["firmware_sha256"] = firmware_sha256(fw)
        version_file = os.path.join(os.path.dirname(fw), "gsp_version.txt")
        if os.path.isfile(version_file):
            with open(version_file, "r", encoding="utf-8", errors="replace") as fh:
                meta["firmware_version"] = "".join(
                    ch for ch in fh.read() if ch.isprintable()
                ).strip()
    return dm


# -- CLI -------------------------------------------------------------------


def _default_image() -> str:
    tool_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    return os.path.join(
        os.path.dirname(os.path.dirname(tool_root)),
        "firmware", "57X-GSP", "570.144", "gsp-570.144.fwimage.bin",
    )


def _cmd_harvest(args) -> int:
    dm = harvest_probe_session(args.session_dir, bias=args.bias)
    print(f"[device-map] harvested {len(dm)} entries from {args.session_dir}", file=sys.stderr)
    if args.out:
        dm.save(args.out)
        print(f"[device-map] wrote {args.out}", file=sys.stderr)
    else:
        sys.stdout.write(dm.to_json())
    return 0


def _cmd_show(args) -> int:
    dm = DeviceMap.load(args.map)
    if args.json:
        sys.stdout.write(dm.to_json())
        return 0
    print(f"map: {args.map}")
    for key in sorted(dm.meta):
        print(f"meta {key}: {dm.meta[key]}")
    print(f"{'address':>20} {'w':>2} {'rw':>3} {'hits':>7} {'first_epc_fw':>14} {'values':>10} source")
    for entry in dm.entries():
        values = ",".join(f"{v:#x}" for v in entry.values) or "-"
        print(
            f"{entry.address:#20x} {entry.width:>2} {entry.kind:>3} {entry.hits:>7} "
            f"{entry.first_epc_fw:#14x} {values:>10} {entry.source}"
        )
    return 0


def _cmd_merge(args) -> int:
    merged = DeviceMap()
    for path in args.maps:
        merged.merge(DeviceMap.load(path))
    merged.save(args.out)
    print(f"[device-map] merged {len(args.maps)} maps -> {args.out} ({len(merged)} entries)")
    return 0


def _run_real_image(args) -> tuple:
    """Load the fwimage into a lenient native Machine and run a fixed budget."""
    from .elf import load_elf
    from sim.machine import Machine

    image = load_elf(args.image)
    bias = args.bias
    end = image.max_vaddr + bias
    need = end - _RAM_BASE + (1 << 20)
    memory_mb = max(256, (need + (1 << 20) - 1) // (1 << 20))
    dm = DeviceMap(meta={"source": "learn", "image": os.path.basename(args.image), "bias": bias})
    machine = Machine(
        memory_mb=memory_mb, device_map=dm, unknown_mmio="lenient", learn_mmio=True, bias=bias
    )
    machine.load_segments(image.flat_load_plan(bias=bias), image.relocated_entry(bias=bias))
    reason, pc, a, b = machine.run(args.instructions)
    return machine, image, dm, reason, pc, a, b


_REASONS = {
    0: "budget", 1: "breakpoint", 2: "watch", 3: "ebreak", 4: "fault", 5: "step", 6: "ecall",
}


def _cmd_demo(args) -> int:
    if not os.path.isfile(args.image):
        print(f"[device-map] image not found: {args.image}", file=sys.stderr)
        return 2
    machine, image, dm, reason, pc, a, b = _run_real_image(args)
    digest = firmware_sha256(args.image)
    print(f"image: {args.image}")
    print(f"sha256: {digest}")
    print(f"elf: entry={image.entry:#x} vaddr=[{image.min_vaddr:#x},{image.max_vaddr:#x})")
    print(f"bias: {args.bias:#x}  relocated entry: {image.relocated_entry(args.bias):#x}")
    print(f"budget: {args.instructions} instructions")
    print(f"stopped: reason={_REASONS.get(reason, reason)} pc={pc:#x} (fw {pc - args.bias:#x})")
    print(f"unknown mmio accesses: {machine.mmio_unknown}")
    for line in machine.mmio_log[: args.first]:
        print(
            f"  first unknown: pc={line['pc']:#x} (fw {line['pc'] - args.bias:#x}) "
            f"addr={line['addr']:#x} size={line['size']} rw={line['rw']}"
        )
    if machine.mmio_log and machine.mmio_log[0]["addr"] == QEMU_REFERENCE["access_addr"]:
        delta = machine.insns - 6  # QEMU traps on the sixth instruction
        print(
            "[device-map] qemu probe first fault: "
            f"fw vaddr {QEMU_REFERENCE['firmware_vaddr']:#x} reading "
            f"{QEMU_REFERENCE['access_addr']:#x} ({QEMU_REFERENCE['desc']})"
        )
        print(
            "[device-map] verdict: first-touch address reproduced; lenient execution "
            f"continued for {machine.insns} instructions ({delta} past the qemu trap)"
        )
    else:
        print("[device-map] verdict: no unknown access at the documented first-touch address")
    if args.out:
        dm.save(args.out)
        print(f"[device-map] learned map ({len(dm)} entries) -> {args.out}")
    return 0


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="python3 -m greeneye.device_map",
        description="Harvest, inspect and merge risc-v-gsp-rm-debugger first-touch device maps.",
    )
    sub = parser.add_subparsers(dest="cmd", required=True)

    harvest = sub.add_parser("harvest", help="harvest a greeneye probe session directory")
    harvest.add_argument("session_dir")
    harvest.add_argument("--out")
    harvest.add_argument("--bias", type=lambda s: int(s, 0), default=None)
    harvest.set_defaults(func=_cmd_harvest)

    show = sub.add_parser("show", help="print a device map")
    show.add_argument("map")
    show.add_argument("--json", action="store_true", help="emit the raw JSON")
    show.set_defaults(func=_cmd_show)

    merge = sub.add_parser("merge", help="merge maps into one (self, then others)")
    merge.add_argument("maps", nargs="+")
    merge.add_argument("--out", required=True)
    merge.set_defaults(func=_cmd_merge)

    demo = sub.add_parser("demo", help="bounded lenient run of a real fwimage")
    demo.add_argument("--image", default=_default_image())
    demo.add_argument("--bias", type=lambda s: int(s, 0), default=_RAM_BASE)
    demo.add_argument("--instructions", type=int, default=200_000)
    demo.add_argument("--first", type=int, default=8)
    demo.add_argument("--out", help="save the learned device map here")
    demo.set_defaults(func=_cmd_demo)

    args = parser.parse_args(argv)
    try:
        return args.func(args)
    except GreenEyeError as exc:
        print(f"[device-map] error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())