"""GSP cycle-0 profile: bounded, deterministic real-firmware run with cycles.

Loads the TU104 570.144 GSP flat image (read-only), maps it the same way the
device-map/probe path does (``bias=0x80000000``, ``greeneye.elf`` PT_LOAD plan,
lenient MMIO from a :class:`greeneye.device_map.DeviceMap`), executes a bounded
instruction budget, and writes a normal ``greeneye-session/1`` directory whose
metadata carries the modeled instruction/cycle totals and the stop reason.

**Fidelity:** this is a bounded lenient first-touch profile, not a boot.  The
cycle figures come from :mod:`sim.cycles` - modeled defaults, not silicon
timings (the simulator's L2 timing is *not* modelled; docs/06-ACCURACY.md,
docs/32-GSP-CYCLE-MODEL.md).  The run is deterministic: two identical
invocations execute the same instructions and report the same cycle total.

CLI::

    python3 -m sim.gsp_profile --instructions 200000
    python3 -m sim.gsp_profile --record dumps/gsp-cycle0 --json

The image is loaded read-only; nothing is written back to it, and no hardware,
GPU, module, guest or network is touched.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from dataclasses import dataclass, field as _field

from greeneye.device_map import DeviceMap, firmware_sha256
from greeneye.elf import load_elf
from greeneye.errors import GreenEyeError
from sim.machine import (
    R_BP,
    R_BUDGET,
    R_EBREAK,
    R_ECALL,
    R_FAULT,
    R_STEP,
    R_WATCH,
)

#: Tool root (the directory containing ``sim/`` and ``greeneye/``).
TOOL_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

#: The TU104 570.144 GSP firmware, as laid out in the workspace tree.
DEFAULT_IMAGE = os.path.normpath(
    os.path.join(
        TOOL_ROOT, "..", "..", "firmware", "57X-GSP", "570.144", "gsp-570.144.fwimage.bin"
    )
)

#: GPU physical addresses are relocated into simulator RAM at this bias
#: (same convention as the probe/device-map path).
DEFAULT_BIAS = 0x80000000

#: Small default budget: enough to pass the documented first-touch fault.
DEFAULT_INSTRUCTIONS = 200_000

#: Machine stop reasons -> stable session labels.
STOP_REASONS = {
    R_BUDGET: "budget",
    R_BP: "breakpoint",
    R_WATCH: "watch",
    R_EBREAK: "ebreak",
    R_FAULT: "fault",
    R_STEP: "step",
    R_ECALL: "ecall",
}


@dataclass
class GspProfileResult:
    """Outcome of one :func:`run_profile` invocation.

    ``ran`` is False (with ``reason`` explaining why) when the image is absent
    and the caller should skip cleanly; ``device_map`` is the in-memory map
    used/learned during the run (never serialised by :meth:`to_dict`).
    """

    ran: bool
    reason: str
    image: str
    image_sha256: str | None = None
    session_dir: str | None = None
    instructions: int = 0
    cycles: int = 0
    cycle_model: str | None = None
    stop_reason: str | None = None
    stop_pc: int | None = None
    stop_pc_fw: int | None = None
    first_unknown_mmio: dict | None = None
    mmio_unknown: int = 0
    device_map: DeviceMap | None = _field(default=None, repr=False)

    def to_dict(self) -> dict:
        data = {
            "schema": "greeneye-gsp-cycle-profile/1",
            "ran": self.ran,
            "reason": self.reason,
            "image": self.image,
            "image_sha256": self.image_sha256,
            "session": self.session_dir,
            "instructions_executed": self.instructions,
            "cycles": self.cycles,
            "cycle_model": self.cycle_model,
            "fidelity": "modeled cycles; L2 (timing) not modelled (docs/06)",
            "stop": {"reason": self.stop_reason, "pc": self.stop_pc, "pc_fw": self.stop_pc_fw},
            "first_unknown_mmio": self.first_unknown_mmio,
            "mmio_unknown": self.mmio_unknown,
        }
        if self.device_map is not None:
            data["device_map_entries"] = len(self.device_map)
        return data


def _fw_pc(pc: int, bias: int) -> int:
    if not bias or pc == 0 or pc < bias:
        return pc
    return pc - bias


def _memory_mb(image, bias: int, ram_base: int = DEFAULT_BIAS) -> int:
    """RAM size that covers the relocated image, as the probe path computes it."""
    end = image.max_vaddr + bias
    need = end - ram_base + (1 << 20)
    return max(256, (need + (1 << 20) - 1) // (1 << 20))


def run_profile(
    *,
    image: str = DEFAULT_IMAGE,
    instructions: int = DEFAULT_INSTRUCTIONS,
    bias: int = DEFAULT_BIAS,
    record=None,
    device_map=None,
    learn_mmio: bool = False,
    name: str = "gsp-cycle-profile",
    cycle_model=None,
) -> GspProfileResult:
    """Run the bounded cycle profile; skip cleanly when the image is absent.

    ``record`` is a session directory (or True for an auto-named one under
    ``dumps/``); the default ``None`` still records, so a session is always
    written.  ``device_map`` may be a :class:`~greeneye.device_map.DeviceMap`,
    a path to one, or None for an empty lenient map.
    """
    if instructions <= 0:
        raise GreenEyeError(f"instructions must be positive, got {instructions}")
    if not os.path.isfile(image):
        return GspProfileResult(ran=False, reason=f"image absent: {image}", image=image)

    image_obj = load_elf(image)
    if not image_obj.is_riscv:
        raise GreenEyeError(f"{image}: not a RISC-V ELF image")
    sha = firmware_sha256(image)

    if isinstance(device_map, (str, os.PathLike)):
        dm = DeviceMap.load(os.fspath(device_map))
    elif device_map is None:
        dm = DeviceMap(
            meta={
                "source": "gsp-profile",
                "firmware": os.path.basename(image),
                "firmware_sha256": sha,
                "bias": bias,
            }
        )
    else:
        dm = device_map

    from greeneye.sim_backend import launch_native  # deferred: keeps sim/ standalone

    session = launch_native(
        kernel=image,
        memory_mb=_memory_mb(image_obj, bias),
        record=True if record is None else record,
        name=name,
        device_map=dm,
        unknown_mmio="lenient",
        learn_mmio=learn_mmio,
        bias=bias,
        cycle_model=cycle_model,
    )
    machine = session.target.machine
    try:
        reason, pc, _a, _b = machine.run(instructions)
        first = machine.mmio_log[0] if machine.mmio_log else None
        model = machine.cycle_model
        stop_reason = STOP_REASONS.get(reason, str(reason))
        stop_pc_fw = _fw_pc(pc, bias)
        meta = {
            "kind": "gsp-cycle-profile",
            "firmware_sha256": sha,
            "bias": bias,
            "instructions_executed": machine.insns,
            "cycles": machine.cycles,
            "cycle_model": model.name if model is not None else None,
            "cycle_fidelity": "modeled cycles, not measured (docs/32-GSP-CYCLE-MODEL.md)",
            "stop_reason": stop_reason,
            "stop_pc": pc,
            "stop_pc_fw": stop_pc_fw,
            "mmio_unknown": machine.mmio_unknown,
            "first_unknown_mmio": first,
        }
        if model is not None:
            meta["cycle_model_table"] = model.describe()
        if getattr(session.recorder, "dir", None) is not None:
            session.recorder.meta.update(meta)
        session.note(
            "gsp_cycle_profile",
            instructions=machine.insns,
            cycles=machine.cycles,
            stop_reason=stop_reason,
            stop_pc_fw=stop_pc_fw,
        )
        session_dir = getattr(session.recorder, "dir", None)
    finally:
        session.close()

    return GspProfileResult(
        ran=True,
        reason="ok",
        image=image,
        image_sha256=sha,
        session_dir=session_dir,
        instructions=machine.insns,
        cycles=machine.cycles,
        cycle_model=model.name if model is not None else None,
        stop_reason=stop_reason,
        stop_pc=pc,
        stop_pc_fw=stop_pc_fw,
        first_unknown_mmio=first,
        mmio_unknown=machine.mmio_unknown,
        device_map=dm,
    )


def _print_result(result: GspProfileResult, out=sys.stdout) -> None:
    print(f"image: {result.image}", file=out)
    print(f"sha256: {result.image_sha256}", file=out)
    print(f"session: {result.session_dir}", file=out)
    print(f"instructions: {result.instructions}", file=out)
    print(f"cycles (modeled): {result.cycles}", file=out)
    print(f"cycle model: {result.cycle_model}", file=out)
    if result.instructions:
        print(f"cycles/instruction: {result.cycles / result.instructions:.3f}", file=out)
    print(
        f"stop: {result.stop_reason} pc={result.stop_pc:#x} (fw {result.stop_pc_fw:#x})",
        file=out,
    )
    print(f"unknown MMIO accesses: {result.mmio_unknown}", file=out)
    if result.first_unknown_mmio:
        first = result.first_unknown_mmio
        print(
            "first unknown MMIO: "
            f"pc={first['pc']:#x} addr={first['addr']:#x} "
            f"size={first['size']} rw={first['rw']}",
            file=out,
        )
    print(
        "note: cycles are modeled (sim/cycles.py defaults); no silicon timing claim",
        file=out,
    )


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        prog="python3 -m sim.gsp_profile",
        description=(
            "Bounded deterministic GSP cycle-0 profile (modeled cycles; no silicon claim). "
            "Loads the real firmware read-only, runs under lenient MMIO and records a session."
        ),
    )
    parser.add_argument("--image", default=DEFAULT_IMAGE, help="flat GSP firmware image")
    parser.add_argument(
        "--instructions", type=int, default=DEFAULT_INSTRUCTIONS, help="instruction budget"
    )
    parser.add_argument("--bias", type=lambda s: int(s, 0), default=DEFAULT_BIAS)
    parser.add_argument("--record", help="session directory (default: auto under dumps/)")
    parser.add_argument("--device-map", help="existing greeneye-device-map/1 JSON to serve MMIO")
    parser.add_argument("--learn", action="store_true", help="learn observed MMIO into the map")
    parser.add_argument("--learn-out", help="save the learned/used map here")
    parser.add_argument("--name", default="gsp-cycle-profile")
    parser.add_argument("--json", action="store_true", help="emit the result as JSON")
    args = parser.parse_args(argv)

    try:
        result = run_profile(
            image=args.image,
            instructions=args.instructions,
            bias=args.bias,
            record=args.record,
            device_map=args.device_map,
            learn_mmio=args.learn,
            name=args.name,
        )
    except GreenEyeError as exc:
        print(f"[gsp-profile] error: {exc}", file=sys.stderr)
        return 1

    if not result.ran:
        print(f"[gsp-profile] skipped: {result.reason}", file=sys.stderr)
        return 2
    if args.learn_out and result.device_map is not None:
        result.device_map.save(args.learn_out)
        print(f"[gsp-profile] device map ({len(result.device_map)} entries) -> {args.learn_out}")
    if args.json:
        print(json.dumps(result.to_dict(), indent=1, sort_keys=True))
    else:
        _print_result(result)
    return 0


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