#!/usr/bin/env python3
"""Bounded deterministic mutation of the demo firmware input surface (R5, H4.3).

Each iteration starts a fresh QEMU session, settles the demo firmware to the
end of its self-test, mutates the mailbox command/args (and optionally a RAM
region), runs it for a bounded budget, and classifies the outcome: normal,
trap (ebreak / control transfer to an invalid PC), fault (first new QEMU
exception, parsed from the guest_errors/int log), no-progress, or
unexpected-result against a small demo oracle. Every attempt is recorded as a
risc-v-gsp-rm-debugger session and a compact finding.

Strategies: invalid command IDs, boundary values, single bit flips of valid
inputs, and buffer length boundaries for the ECHO/FILL/SUM commands.

Safety: the harness launches its own QEMU on the bundled (or explicitly
allowed) demo/re-hosted ELF only. It never attaches to hardware, a guest or a
kernel module, and it never uses the native simulator or an existing gdbstub.

Run from the tool root:

    python3 harness/mutation_harness.py --out dumps/mutation --iterations 64 --seed 7
"""
from __future__ import annotations

import argparse
import json
import os
import random
import shutil
import struct
import sys
import tempfile
from collections import Counter
from dataclasses import dataclass

try:
    from .common import (
        DEMO_ELF,
        LOAD_ADDR,
        MAX_LOG_BYTES,
        ensure_dir,
        load_target,
        memory_mb_for,
        open_flat_session,
        qemu_available,
        read_faults_window,
        sha256_file,
    )
except ImportError:  # direct script execution (python3 harness/mutation_harness.py)
    from common import (  # type: ignore[no-redef]
        DEMO_ELF,
        LOAD_ADDR,
        MAX_LOG_BYTES,
        ensure_dir,
        load_target,
        memory_mb_for,
        open_flat_session,
        qemu_available,
        read_faults_window,
        sha256_file,
    )

from greeneye.errors import GreenEyeError  # noqa: E402

FORMAT = "greeneye-mutation-findings/1"
MASK64 = (1 << 64) - 1
VALID_CMDS = (1, 2, 3, 4, 5, 6)
INVALID_CMDS = (0, 7, 8, 15, 24, 25, 31, 32, 63, 0x100, 0xFFFF, 0xFFFFFFFF, MASK64)
BOUNDARY_VALUES = (
    0,
    1,
    2,
    3,
    4,
    7,
    8,
    16,
    32,
    63,
    64,
    65,
    128,
    0xFF,
    0x100,
    0x101,
    1024,
    0x7FFFFFFF,
    0x80000000,
    0x40000000,
    0xFFFFFFFF,
    0x8000000000000000,
    MASK64,
)
BUFFER_LENGTHS = (0, 1, 2, 31, 32, 33, 255, 256, 257, 512, 4096, 0x10000, 0x10000000, MASK64)
STRATEGIES = ("cmd-invalid", "arg-boundary", "bit-flip", "buffer-length")

CTRL_MAGIC = 0
CTRL_VERSION = 4
CTRL_STATE = 8
CTRL_RPC_COUNT = 16
CTRL_LAST_RPC = 24
CTRL_MAILBOX = 32
CTRL_ARG0 = 40
CTRL_ARG1 = 48
CTRL_ARG2 = 56
CTRL_RESULT = 64
CTRL_HANDLER_COUNT = 328
CTRL_BUF_OFFSET = 72  # char buf[256]
CTRL_SIZE = 400
SETTLE_MARKER = "SELFTEST DONE"
SETTLE_RETRIES = 3


@dataclass
class Mutation:
    iteration: int
    strategy: str
    cmd: int = 0
    arg0: int | str = 0
    arg1: int = 0
    arg2: int = 0
    region: dict | None = None

    def to_json(self) -> dict:
        return {
            "cmd": self.cmd,
            "arg0": self.arg0,
            "arg1": self.arg1,
            "arg2": self.arg2,
            "region": self.region,
        }


@dataclass
class HarnessOptions:
    kernel: str
    out: str
    iterations: int = 50
    seed: int = 1
    timeout: float = 0.25
    settle: float = 0.2
    phase_a: float = 0.05
    memory_region: tuple[int, int] | None = None
    allow_foreign_elf: bool = False
    max_log_bytes: int = MAX_LOG_BYTES
    kernel_sha256: str = ""
    is_demo: bool = False


# ---------------------------------------------------------------------------
# mutation generation
# ---------------------------------------------------------------------------


def generate_mutations(iterations: int, seed: int, region: tuple[int, int] | None = None) -> list[Mutation]:
    """Deterministic strategy rotation; per-iteration RNG seeded by (seed, index)."""
    strategies = STRATEGIES + (("mem-region",) if region else ())
    mutations: list[Mutation] = []
    for index in range(iterations):
        # Integer seed: deterministic across Python versions/processes.
        rng = random.Random(seed * 1_000_003 + index)
        strategy = strategies[index % len(strategies)]
        mutation = Mutation(iteration=index, strategy=strategy)
        if strategy == "cmd-invalid":
            mutation.cmd = rng.choice(INVALID_CMDS)
        elif strategy == "arg-boundary":
            mutation.cmd = rng.choice(VALID_CMDS)
            args = [0, 0, 0]
            args[rng.randrange(3)] = rng.choice(BOUNDARY_VALUES)
            mutation.arg0, mutation.arg1, mutation.arg2 = args
        elif strategy == "bit-flip":
            values = {
                "cmd": rng.choice(VALID_CMDS),
                "arg0": rng.choice((0, 1, 0x41, 0xFF)),
                "arg1": rng.choice((0, 1, 32)),
                "arg2": 0,
            }
            target = rng.choice(("cmd", "arg0", "arg1", "arg2"))
            values[target] = (values[target] ^ (1 << rng.randrange(64))) & MASK64
            mutation.cmd = values["cmd"]
            mutation.arg0, mutation.arg1, mutation.arg2 = values["arg0"], values["arg1"], values["arg2"]
        elif strategy == "buffer-length":
            mutation.cmd = rng.choice((2, 3, 4))
            mutation.arg0 = "ctrl_buf" if mutation.cmd in (2, 4) else rng.choice((0, 0x41))
            mutation.arg1 = rng.choice(BUFFER_LENGTHS)
        elif strategy == "mem-region":
            addr, size = region  # type: ignore[misc]
            mutation.region = {
                "addr": addr,
                "size": size,
                "offset": rng.randrange(size),
                "xor": 1 << rng.randrange(8),
                "old": None,
                "new": None,
            }
        mutations.append(mutation)
    return mutations


def resolve_arg0(mutation: Mutation, ctrl: int) -> int:
    if mutation.arg0 == "ctrl_buf":
        return ctrl + CTRL_BUF_OFFSET
    return int(mutation.arg0)


# ---------------------------------------------------------------------------
# demo oracle (documented subset only; everything else is not judged)
# ---------------------------------------------------------------------------


def expected_for(mutation: Mutation, before: dict, ctrl: int) -> dict:
    cmd = mutation.cmd
    if cmd == 0:
        return {"mode": "same"}
    if cmd == 1:
        return {"mode": "value", "value": (resolve_arg0(mutation, ctrl) ^ 0xBEEF) & MASK64}
    if cmd == 3:
        return {"mode": "value", "value": min(mutation.arg1, 256)}
    if cmd == 5:
        return {"mode": "one_of", "values": [0x10, 0x20, 0x30, 0x40]}
    if 16 <= cmd <= 21:
        return {"mode": None, "reason": "injected handler slot result is service-dependent"}
    if cmd in (2, 4, 6):
        return {"mode": None, "reason": "memory/trap-dependent handler"}
    return {"mode": "value", "value": (~cmd) & MASK64}


def expected_ok(expectation: dict, mutation: Mutation, before: dict, after: dict) -> bool | None:
    mode = expectation.get("mode")
    if mode is None:
        return None
    if mode == "same":
        return after["result"] == before["result"] and after["last_rpc"] == before["last_rpc"]
    if mode == "value":
        return after["result"] == expectation["value"] and after["last_rpc"] == mutation.cmd
    if mode == "one_of":
        return after["result"] in expectation["values"] and after["last_rpc"] == mutation.cmd
    return None


# ---------------------------------------------------------------------------
# iteration execution
# ---------------------------------------------------------------------------


def read_ctrl_state(session, ctrl: int) -> dict:
    raw = session.memory.read(ctrl, CTRL_SIZE)
    u32 = lambda offset: struct.unpack_from("<I", raw, offset)[0]  # noqa: E731
    u64 = lambda offset: struct.unpack_from("<Q", raw, offset)[0]  # noqa: E731
    return {
        "magic": u32(CTRL_MAGIC),
        "version": u32(CTRL_VERSION),
        "state": u32(CTRL_STATE),
        "rpc_count": u64(CTRL_RPC_COUNT),
        "last_rpc": u64(CTRL_LAST_RPC),
        "result": u64(CTRL_RESULT),
        "handler_count": u64(CTRL_HANDLER_COUNT),
    }


def apply_mutation(session, ctrl: int, mutation: Mutation) -> None:
    if mutation.region:
        region = mutation.region
        data = bytearray(session.memory.read(region["addr"], region["size"]))
        region["old"] = data[region["offset"]]
        data[region["offset"]] ^= region["xor"]
        region["new"] = data[region["offset"]]
        session.memory.write(region["addr"], bytes(data))
    session.memory.write_u64(ctrl + CTRL_MAILBOX, mutation.cmd)
    session.memory.write_u64(ctrl + CTRL_ARG0, resolve_arg0(mutation, ctrl))
    session.memory.write_u64(ctrl + CTRL_ARG1, mutation.arg1)
    session.memory.write_u64(ctrl + CTRL_ARG2, mutation.arg2)


def _compact_fault(fault: dict) -> dict:
    return {
        "cause": fault.get("cause"),
        "desc": fault.get("desc"),
        "epc": fault.get("epc"),
        "tval": fault.get("tval"),
        "access": fault.get("access"),
    }


def classify(
    mutation: Mutation,
    before: dict,
    after: dict,
    faults: list[dict],
    trap_loop: bool,
    exit_seen: bool,
    settle_ok: bool,
    ctrl: int,
) -> tuple[str, dict]:
    delta = after["rpc_count"] - before["rpc_count"]
    detail = {
        "observed": {
            "result": after["result"],
            "last_rpc": after["last_rpc"],
            "rpc_count_delta": delta,
        },
        "expected": None,
        "expected_ok": None,
    }
    if not settle_ok:
        detail["reason"] = f"{SETTLE_MARKER!r} not observed within the settle budget"
        return "settle-failed", detail
    if exit_seen:
        return "exit", detail
    if faults:
        fault = faults[0]
        detail["fault"] = _compact_fault(fault)
        return ("trap" if fault.get("cause") == 3 else "fault"), detail
    if trap_loop:
        return "trap", detail
    if mutation.cmd != 0 and delta <= 0:
        return "no-progress", detail
    expectation = expected_for(mutation, before, ctrl)
    detail["expected"] = expectation
    verdict = expected_ok(expectation, mutation, before, after)
    detail["expected_ok"] = verdict
    if verdict is False:
        return "unexpected-result", detail
    return "ok", detail


def _settle(session, fs, options: HarnessOptions) -> bool:
    for _ in range(SETTLE_RETRIES):
        session.run(duration=options.settle, max_stops=1)
        if SETTLE_MARKER in fs.target.read_uart():
            return True
    return False


def run_iteration(mutation: Mutation, options: HarnessOptions, image) -> dict:
    session_rel = os.path.join("sessions", f"iter-{mutation.iteration:04d}")
    finding = {
        "iteration": mutation.iteration,
        "strategy": mutation.strategy,
        "mutation": mutation.to_json(),
        "outcome": "error",
        "detail": {},
        "session": session_rel,
    }
    workdir = tempfile.mkdtemp(prefix="greeneye-mut-")
    try:
        with open_flat_session(
            image,
            options.kernel,
            record=os.path.join(options.out, session_rel),
            name=f"mut-{mutation.iteration:04d}",
            symbols=None,
            log_faults=True,
            workdir=workdir,
        ) as fs:
            session = fs.session
            ctrl_symbol = session.symbols.get("gsp_ctrl")
            if ctrl_symbol is None:
                raise GreenEyeError("gsp_ctrl symbol not found: the mutation harness targets the demo control block")
            ctrl = ctrl_symbol.addr
            settle_ok = _settle(session, fs, options)
            before = read_ctrl_state(session, ctrl)
            faults: list[dict] = []
            trap_loop = False
            exit_seen = False
            after = before
            if settle_ok:
                log_offset = os.path.getsize(fs.target.qemu_log) if os.path.exists(fs.target.qemu_log) else 0
                apply_mutation(session, ctrl, mutation)
                phase_a = min(options.phase_a, options.timeout)
                session.run(duration=phase_a, max_stops=1)
                faults, _capped = read_faults_window(fs.target.qemu_log, log_offset, options.max_log_bytes)
                trap_loop = session.pc < 0x100
                after = read_ctrl_state(session, ctrl)
                # Slow/hung handlers get one confirmation window before we call
                # it no-progress; faults are still picked up if they appear.
                if not faults and not trap_loop and mutation.cmd != 0 and after["rpc_count"] == before["rpc_count"]:
                    session.run(duration=max(0.01, options.timeout - phase_a), max_stops=1)
                    more, _capped = read_faults_window(fs.target.qemu_log, log_offset, options.max_log_bytes)
                    faults = more or faults
                    trap_loop = session.pc < 0x100
                    after = read_ctrl_state(session, ctrl)
                exit_seen = not fs.target.alive()
            outcome, detail = classify(mutation, before, after, faults, trap_loop, exit_seen, settle_ok, ctrl)
            finding["outcome"] = outcome
            finding["detail"] = detail
    except Exception as exc:  # one bad iteration must not end the run
        finding["outcome"] = "error"
        finding["detail"] = {"reason": f"{type(exc).__name__}: {exc}"}
    finally:
        shutil.rmtree(workdir, ignore_errors=True)
    return finding


# ---------------------------------------------------------------------------
# run + report
# ---------------------------------------------------------------------------


def run_harness(options: HarnessOptions) -> list[dict]:
    if not qemu_available():
        raise GreenEyeError("qemu-system-riscv64 not found; install qemu-system-misc or set GREENEYE_QEMU")
    if not os.path.isfile(options.kernel):
        raise GreenEyeError(f"kernel not found: {options.kernel}")
    image = load_target(options.kernel)
    options.kernel_sha256 = sha256_file(options.kernel)
    demo_sha = sha256_file(DEMO_ELF) if os.path.isfile(DEMO_ELF) else None
    options.is_demo = options.kernel_sha256 == demo_sha or os.path.abspath(options.kernel) == os.path.abspath(DEMO_ELF)
    if not options.is_demo and not options.allow_foreign_elf:
        raise GreenEyeError(
            "mutation is restricted to the bundled demo firmware; "
            "pass --allow-foreign-elf to mutate another re-hosted ELF under QEMU"
        )
    if options.memory_region:
        addr, size = options.memory_region
        limit = LOAD_ADDR + memory_mb_for(image) * 1024 * 1024
        if size <= 0 or addr < LOAD_ADDR or addr + size > limit:
            raise GreenEyeError(
                f"--memory-region {addr:#x}:{size:#x} is outside QEMU RAM "
                f"[{LOAD_ADDR:#x}, {limit:#x})"
            )
    ensure_dir(options.out)
    findings: list[dict] = []
    mutations = generate_mutations(options.iterations, options.seed, options.memory_region)
    findings_path = os.path.join(options.out, "findings.jsonl")
    with open(findings_path, "w", encoding="utf-8") as fh:
        for mutation in mutations:
            finding = run_iteration(mutation, options, image)
            findings.append(finding)
            fh.write(json.dumps(finding, sort_keys=True) + "\n")
            fh.flush()
    write_summary(options, findings)
    return findings


def write_summary(options: HarnessOptions, findings: list[dict]) -> str:
    classes = Counter(finding["outcome"] for finding in findings)
    strategies = Counter(finding["strategy"] for finding in findings)
    lines = [
        "# risc-v-gsp-rm-debugger mutation findings",
        "",
        "> Status: implemented, documented and tested (audit: docs/27-TOOL-AUDIT.md; value demos: docs/25-VALUE-VALIDATION.md). Generated by `harness/mutation_harness.py`.",
        "",
        f"* format: `{FORMAT}`",
        f"* kernel: `{options.kernel}`",
        f"* kernel sha256: `{options.kernel_sha256}`",
        f"* demo restriction: {'demo' if options.is_demo else 'foreign ELF (explicitly allowed)'}",
        f"* seed: {options.seed}, iterations: {options.iterations}",
        f"* timeout: {options.timeout}s, settle: {options.settle}s",
        f"* memory region: {_fmt_region(options.memory_region)}",
        "",
        "## Outcome classes",
        "",
    ]
    lines += _table(
        ["outcome", "count", "first iteration"],
        [
            [outcome, classes[outcome], next(f["iteration"] for f in findings if f["outcome"] == outcome)]
            for outcome in sorted(classes)
        ],
    )
    lines += ["", "## Strategy coverage", ""]
    lines += _table(
        ["strategy", "iterations", "outcomes"],
        [
            [
                strategy,
                strategies[strategy],
                ", ".join(
                    f"{outcome}={count}"
                    for outcome, count in sorted(Counter(f["outcome"] for f in findings if f["strategy"] == strategy).items())
                ),
            ]
            for strategy in sorted(strategies)
        ],
    )
    notable = [finding for finding in findings if finding["outcome"] != "ok"][:20]
    lines += ["", "## Notable findings (first 20 non-ok)", ""]
    if not notable:
        lines += ["_all iterations completed normally_", ""]
    for finding in notable:
        lines += [
            f"* iteration {finding['iteration']} `{finding['strategy']}` "
            f"cmd={finding['mutation']['cmd']:#x} -> **{finding['outcome']}** "
            f"{_short_json(finding['detail'])}",
        ]
    lines += ["", "## Reproduce", "", "```bash", _reproduce_command(options), "```", ""]
    path = os.path.join(options.out, "summary.md")
    with open(path, "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines))
    return path


def _reproduce_command(options: HarnessOptions) -> str:
    region = f" --memory-region {_fmt_region(options.memory_region)}" if options.memory_region else ""
    foreign = " --allow-foreign-elf" if not options.is_demo else ""
    return (
        f"python3 harness/mutation_harness.py --kernel {options.kernel} --out {options.out} "
        f"--iterations {options.iterations} --seed {options.seed} --timeout {options.timeout} "
        f"--settle {options.settle}{region}{foreign}"
    )


def _fmt_region(region: tuple[int, int] | None) -> str:
    return "-" if not region else f"{region[0]:#x}:{region[1]:#x}"


def _short_json(value, limit: int = 120) -> str:
    text = json.dumps(value, sort_keys=True)
    return f"`{text}`" if len(text) <= limit else f"`{text[:limit - 1]}...`"


def _table(headers: list[str], rows: list[list]) -> list[str]:
    lines = ["| " + " | ".join(headers) + " |", "|" + "|".join("---" for _ in headers) + "|"]
    for row in rows:
        lines.append("| " + " | ".join(str(cell) for cell in row) + " |")
    return lines


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


def parse_region(text: str) -> tuple[int, int]:
    try:
        addr_text, size_text = text.split(":", 1)
        addr = int(addr_text, 0)
        size = int(size_text, 0)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(f"expected ADDR:SIZE, got {text!r}") from exc
    return addr, size


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="mutation_harness", description=__doc__.splitlines()[0])
    parser.add_argument("--kernel", default=DEMO_ELF, help="demo/re-hosted ELF to mutate (default: bundled demo)")
    parser.add_argument("--out", required=True, help="output directory for findings and recorded sessions")
    parser.add_argument("--iterations", type=int, default=50, help="mutation count (1..10000)")
    parser.add_argument("--seed", type=int, default=1, help="deterministic RNG seed")
    parser.add_argument("--timeout", type=float, default=0.25, help="total observation budget per mutation (seconds)")
    parser.add_argument("--settle", type=float, default=0.2, help="settle budget per attempt before mutating (seconds)")
    parser.add_argument("--phase-a", type=float, default=0.05, help="first observation window before confirmation")
    parser.add_argument("--memory-region", type=parse_region, help="optional RAM region to mutate, ADDR:SIZE")
    parser.add_argument("--allow-foreign-elf", action="store_true", help="permit a non-demo ELF (QEMU only)")
    parser.add_argument("--max-log-bytes", type=int, default=MAX_LOG_BYTES, help="QEMU log window per fault check")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if not 1 <= args.iterations <= 10000:
        parser.error("--iterations must be in 1..10000")
    if not 0.01 <= args.timeout <= 10.0:
        parser.error("--timeout must be in 0.01..10.0")
    if not 0.0 <= args.settle <= 10.0:
        parser.error("--settle must be in 0..10.0")
    options = HarnessOptions(
        kernel=os.path.abspath(args.kernel),
        out=os.path.abspath(args.out),
        iterations=args.iterations,
        seed=args.seed,
        timeout=args.timeout,
        settle=args.settle,
        phase_a=args.phase_a,
        memory_region=args.memory_region,
        allow_foreign_elf=args.allow_foreign_elf,
        max_log_bytes=args.max_log_bytes,
    )
    try:
        findings = run_harness(options)
    except GreenEyeError as exc:
        print(f"mutation harness error: {exc}", file=sys.stderr)
        return 2
    classes = Counter(finding["outcome"] for finding in findings)
    print(f"[mutation] {len(findings)} iteration(s), out={options.out}")
    for outcome in sorted(classes):
        print(f"[mutation] {outcome}: {classes[outcome]}")
    print(f"[mutation] wrote {os.path.join(options.out, 'findings.jsonl')}")
    print(f"[mutation] wrote {os.path.join(options.out, 'summary.md')}")
    return 0


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())