#!/usr/bin/env python3
"""Surgical rehost: run a bounded function of a firmware image on the native emulator.

The standalone simulator boots the real 570.144 image to the message-loop port
wait (B3).  Code *after* that wall (the RM handlers, e.g. the GR promote/commit
chain) is only reachable when host traffic exists, which is the B4/B5 rung.
This harness provides the offline alternative the workflow review ranks first
(docs/39 U2): start the machine at a chosen function entry with a crafted
register/memory state, run forward with a bounded budget, and stop on return,
fault, unhandled ecall or budget.  Every run records a session and a report;
snapshots before/after a run show what the function changed.

Honesty rules:

* This tool executes firmware code; it never touches a GPU, kernel module,
  mdev, guest or network.  It is development-only.
* A rehosted run is **not** a boot: the function runs with only the state the
  caller supplied.  The report says so, and any fault (first missing state) is
  recorded as evidence, never papered over.
* No timing/cycle claim; fidelity is L0/L1 (docs/06).
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import time
from typing import Any

TOOL_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if TOOL_ROOT not in sys.path:
    sys.path.insert(0, TOOL_ROOT)

from greeneye.elf import load_elf  # noqa: E402
from greeneye.errors import GreenEyeError  # noqa: E402
from greeneye.sim_backend import launch_native  # noqa: E402
from sim.devices import load_profile, registry_from_profile  # noqa: E402
from sim.gsp_boot import _apply_boot_shim  # noqa: E402
from sim.libos import LibosSyscalls  # noqa: E402
from sim.machine import R_BP, R_ECALL, R_FAULT, R_WATCH  # noqa: E402

#: Schema tag of the rehost report.
REHOST_FORMAT = "greeneye-rehost/1"
#: Schema tag of the demo result.
DEMO_FORMAT = "greeneye-rehost-demo/1"

DEFAULT_IMAGE = os.path.join(
    TOOL_ROOT, "..", "..", "firmware", "57X-GSP", "570.144", "gsp-570.144.fwimage.bin"
)
DEFAULT_PROFILE = os.path.join(TOOL_ROOT, "sim", "profiles", "tu104-570.144.json")
DEMO_ELF = os.path.join(TOOL_ROOT, "firmware", "uproc_demo", "build", "uproc_demo.elf")

#: Wall-clock guard: a rehost that spins in a fault loop must not hang the tool.
DEFAULT_TIMEOUT = 60.0
#: Instruction budget per default run.
DEFAULT_BUDGET = 5_000_000
#: Chunk size between stop checks (keeps the loop responsive to timeouts).
RUN_CHUNK = 250_000
#: Sentinel instruction (``ebreak``) written at the return trampoline.
EBREAK = 0x00100073


class RehostError(GreenEyeError):
    """A rehost input, image or argument is invalid."""


def _parse_u64(value: Any) -> int:
    if isinstance(value, bool):
        return int(value)
    if isinstance(value, int):
        return value & 0xFFFF_FFFF_FFFF_FFFF
    if isinstance(value, str):
        text = value.strip()
        if not text:
            raise RehostError("empty integer")
        return int(text, 0) & 0xFFFF_FFFF_FFFF_FFFF
    raise RehostError(f"cannot parse integer from {value!r}")


def parse_hex_bytes(spec: str) -> bytes:
    """Parse ``aabbcc`` (whitespace tolerated) into bytes."""
    text = "".join(spec.split())
    if len(text) % 2:
        raise RehostError(f"hex payload must have an even number of digits: {spec!r}")
    try:
        return bytes.fromhex(text)
    except ValueError as exc:
        raise RehostError(f"bad hex payload {spec!r}: {exc}") from exc


def parse_mem_spec(spec: str) -> tuple[int, bytes]:
    """Parse ``<addr>=<hex>`` into ``(address, bytes)``."""
    if "=" not in spec:
        raise RehostError(f"memory spec must be <addr>=<hex>: {spec!r}")
    addr_text, _, data_text = spec.partition("=")
    return _parse_u64(addr_text), parse_hex_bytes(data_text)


def parse_snap_spec(spec: str) -> tuple[int, int]:
    """Parse ``<addr>:<size>`` into ``(address, size)``."""
    addr_text, sep, size_text = spec.partition(":")
    if not sep:
        raise RehostError(f"snapshot spec must be <addr>:<size>: {spec!r}")
    size = _parse_u64(size_text)
    if size <= 0:
        raise RehostError(f"snapshot size must be positive: {spec!r}")
    return _parse_u64(addr_text), size


def _sha256_file(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()


def _region_diff(before: bytes, after: bytes) -> dict:
    changed = 0
    first = None
    for offset, (a, b) in enumerate(zip(before, after)):
        if a != b:
            changed += 1
            if first is None:
                first = offset
    return {
        "changed_bytes": changed,
        "first_diff_offset": first,
        "before_sha256": hashlib.sha256(before).hexdigest(),
        "after_sha256": hashlib.sha256(after).hexdigest(),
    }


def _registers(machine) -> dict:
    regs = {f"x{i}": machine.x[i] for i in range(32) if i and machine.x[i]}
    regs["pc"] = machine.pc
    return regs


def rehost(
    *,
    image: str,
    entry: int | str,
    profile_path: str | None = None,
    args: tuple = (),
    mem: tuple = (),
    snap: tuple = (),
    budget: int = DEFAULT_BUDGET,
    timeout: float = DEFAULT_TIMEOUT,
    out: str | None = None,
    record=None,
    name: str = "rehost",
    apply_shim: bool = True,
    unknown_mmio: str = "lenient",
) -> dict:
    """Run one bounded function call on the native emulator; return a report."""
    if not os.path.isfile(image):
        raise RehostError(f"image not found: {image}")
    if budget <= 0:
        raise RehostError(f"budget must be positive, got {budget}")
    image_obj = load_elf(image)
    if not image_obj.is_riscv:
        raise RehostError(f"{image}: not a RISC-V ELF image")

    profile = load_profile(profile_path) if profile_path else None
    load = (profile or {}).get("load", {})
    ram_base = _parse_u64(load["ram_base"]) if "ram_base" in load else 0x80000000
    memory_mb = int(load.get("memory_mb", 256))

    session = launch_native(
        kernel=image,
        memory_mb=memory_mb,
        ram_base=ram_base,
        record=record if record is not None else bool(out),
        name=name,
        device_map=None,
        unknown_mmio=unknown_mmio,
        bias=0,
    )
    machine = session.target.machine
    try:
        if profile is not None:
            machine.devices = registry_from_profile(profile, machine=machine)
            machine.translator = getattr(machine.devices, "dma", None)
            libos = LibosSyscalls()
            machine.syscall_handler = libos.handle
            if apply_shim:
                shim_meta = {"enabled": True, "assumptions": []}
                shim = dict(profile.get("boot_shim", {}))
                shim["_image_path"] = image
                _apply_boot_shim(machine, shim, image_obj, shim_meta)

        # Resolve a symbol name through the ELF symbol table when possible.
        if isinstance(entry, str):
            try:
                entry_addr = session.symbols.lookup(entry) if session.symbols else None
            except GreenEyeError:
                entry_addr = None
            if entry_addr is None:
                raise RehostError(f"symbol {entry!r} not found in {os.path.basename(image)}")
        else:
            entry_addr = _parse_u64(entry)

        for addr, data in mem:
            machine.write_mem(addr, data)

        snapshots = []
        for addr, size in snap:
            before = machine.read_mem(addr, size)
            snapshots.append({"addr": addr, "size": size, "before": before})

        # Return trampoline well above the image, inside the RAM arena.
        sentinel = (machine.ram_base + machine.ram_size - 0x100) & ~0xF
        stack = (sentinel - 0x2000) & ~0xF
        trampoline = EBREAK.to_bytes(4, "little")
        machine.write_mem(sentinel, trampoline)
        machine.write_mem(sentinel + 4, trampoline)
        for index, value in enumerate(tuple(args)[:8]):
            machine.x[10 + index] = _parse_u64(value)
        machine.x[1] = sentinel  # ra
        machine.x[2] = stack  # sp
        machine.pc = entry_addr

        entered = _registers(machine)
        started = time.monotonic()
        status = "budget"
        fault: dict | None = None
        stop_pc = machine.pc
        remaining = budget
        while remaining > 0:
            if time.monotonic() - started > timeout:
                status = "timeout"
                break
            chunk = min(RUN_CHUNK, remaining)
            reason, stop_pc, a, b = machine.run(chunk, breakpoints={sentinel})
            remaining -= chunk
            if reason == R_BP:
                status = "return"
                break
            if reason == R_FAULT:
                status = "fault"
                fault = {"cause": a, "tval": f"{b:#x}"}
                break
            if reason == R_ECALL:
                handled, result = machine.syscall_handler(machine, stop_pc) if machine.syscall_handler else (False, None)
                if not handled:
                    status = "ecall"
                    break
                if result is not None:
                    machine.x[10] = result & 0xFFFF_FFFF_FFFF_FFFF
                continue
            if reason == R_WATCH:
                continue  # watchpoints recorded by the caller's own tooling
            # R_BUDGET: keep going until the total budget is spent.
        if status == "budget" and remaining <= 0:
            status = "budget"

        report: dict[str, Any] = {
            "format": REHOST_FORMAT,
            "image": os.path.abspath(image),
            "image_sha256": _sha256_file(image),
            "profile": os.path.basename(profile_path) if profile_path else None,
            "entry": f"{entry_addr:#x}",
            "args": [_parse_u64(v) for v in tuple(args)[:8]],
            "status": status,
            "stop_pc": f"{stop_pc:#x}",
            "fault": fault,
            "retval": machine.x[10],
            "instructions": 0 if machine.insns is None else machine.insns,
            "wall_seconds": round(time.monotonic() - started, 4),
            "registers_entered": {k: f"{v:#x}" if isinstance(v, int) else v for k, v in entered.items()},
            "registers_exited": {k: f"{v:#x}" if isinstance(v, int) else v for k, v in _registers(machine).items()},
            "mmio_unknown": machine.mmio_unknown,
            "snapshots": [],
            "limits": [
                "Not a boot: the function ran with only the supplied state.",
                "No timing/cycle claim (L0/L1 fidelity, docs/06).",
                "A fault is the first missing state observed, not proof of a product defect.",
            ],
        }
        for item in snapshots:
            after = machine.read_mem(item["addr"], item["size"])
            report["snapshots"].append(
                {
                    "addr": f"{item['addr']:#x}",
                    "size": item["size"],
                    **_region_diff(item["before"], after),
                }
            )
        if getattr(session.recorder, "dir", None) is not None:
            session.note("rehost", entry=f"{entry_addr:#x}", status=status, stop_pc=f"{stop_pc:#x}")
            report["session"] = session.recorder.dir
        if out:
            os.makedirs(out, exist_ok=True)
            path = os.path.join(out, "rehost-report.json")
            with open(path, "w", encoding="utf-8") as fh:
                json.dump(report, fh, indent=2, sort_keys=True)
                fh.write("\n")
            report["report_json"] = os.path.abspath(path)
        return report
    finally:
        session.close()


def run_demo(out: str | None = None) -> dict:
    """Rehost a known demo function to prove the loop (synthetic, deterministic)."""
    if not os.path.isfile(DEMO_ELF):
        raise RehostError(
            f"demo firmware not built: {DEMO_ELF} (run `make demo` in the tool dir)"
        )
    buffer_addr, payload = 0x80012000, bytes([1, 2, 3, 4])
    expected = sum(payload)
    first = rehost(
        image=DEMO_ELF,
        entry="rpc_sum",
        args=(buffer_addr, len(payload)),
        mem=((buffer_addr, payload),),
        snap=(),
        budget=100_000,
        timeout=20.0,
        record=False,
        name="rehost-demo",
    )
    second = rehost(
        image=DEMO_ELF,
        entry="rpc_sum",
        args=(buffer_addr, len(payload)),
        mem=((buffer_addr, payload),),
        snap=(),
        budget=100_000,
        timeout=20.0,
        record=False,
        name="rehost-demo",
    )
    ok = (
        first["status"] == "return"
        and first["retval"] == expected
        and first["status"] == second["status"]
        and first["retval"] == second["retval"]
    )
    payload_result = {
        "format": DEMO_FORMAT,
        "ok": ok,
        "expected": expected,
        "first": {
            "status": first["status"],
            "retval": first["retval"],
            "wall_seconds": first["wall_seconds"],
        },
        "second": {
            "status": second["status"],
            "retval": second["retval"],
        },
    }
    if out:
        os.makedirs(out, exist_ok=True)
        path = os.path.join(out, "rehost-demo.json")
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(payload_result, fh, indent=2, sort_keys=True)
            fh.write("\n")
        payload_result["artifacts"] = {"demo_json": os.path.abspath(path)}
    return payload_result


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="rehost.py",
        description=(
            "Run one bounded firmware function on the native emulator "
            "(surgical rehost). Offline only; never touches hardware."
        ),
    )
    sub = parser.add_subparsers(dest="command")
    run_cmd = sub.add_parser("run", help="rehost a function from an image")
    run_cmd.add_argument("--image", default=DEFAULT_IMAGE, help="firmware ELF")
    run_cmd.add_argument("--profile", default=DEFAULT_PROFILE, help="device profile JSON")
    run_cmd.add_argument("--entry", required=True, help="function entry (hex address or symbol)")
    run_cmd.add_argument("--arg", action="append", default=[], help="a0..a7 value (repeatable)")
    run_cmd.add_argument("--mem", action="append", default=[], help="<addr>=<hex> pre-write (repeatable)")
    run_cmd.add_argument("--snap", action="append", default=[], help="<addr>:<size> snapshot (repeatable)")
    run_cmd.add_argument("--budget", type=int, default=DEFAULT_BUDGET)
    run_cmd.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT)
    run_cmd.add_argument("--no-shim", action="store_true", help="do not apply the profile boot shim")
    run_cmd.add_argument("--out", help="output directory (report + session)")
    demo_cmd = sub.add_parser("demo", help="synthetic self-test on the demo firmware")
    demo_cmd.add_argument("--out", help="output directory")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        if args.command == "run":
            report = rehost(
                image=args.image,
                profile_path=args.profile or None,
                entry=args.entry,
                args=tuple(args.arg),
                mem=tuple(parse_mem_spec(spec) for spec in args.mem),
                snap=tuple(parse_snap_spec(spec) for spec in args.snap),
                budget=args.budget,
                timeout=args.timeout,
                out=args.out,
                apply_shim=not args.no_shim,
            )
        elif args.command == "demo":
            report = run_demo(args.out)
        else:
            parser.print_help()
            return 2
    except RehostError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2
    summary_keys = (
        "format", "ok", "entry", "status", "stop_pc", "retval", "fault",
        "instructions", "wall_seconds", "report_json", "session", "artifacts",
    )
    print(json.dumps({k: report[k] for k in summary_keys if k in report}, indent=2, sort_keys=True))
    if args.command == "demo" and not report.get("ok"):
        return 3
    return 0


if __name__ == "__main__":
    raise SystemExit(main())