#!/usr/bin/env python3
"""Build/run offline CPU/MMIO LBA trace configs for GDR-8050L proof packs.

The v169 locator produces candidate static findings.  This tool turns those
candidate findings into direct mn103s_emu.py configs with focused breakpoints and
watchpoints, then optionally runs the emulator and validates the resulting trace.
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from copy import deepcopy
from pathlib import Path
from typing import Any

from hlds_profiles import load_json_or_jsonc, write_json
from hlds_lba_trace_validator import validate_lba_trace, write_md as write_lba_md


def parse_int(v: Any, default: int = 0) -> int:
    if v is None or v == "":
        return default
    if isinstance(v, int):
        return v
    return int(str(v).strip(), 0)


def hx(data: bytes | list[int]) -> str:
    return " ".join(f"{int(b) & 0xFF:02X}" for b in data)



def require_existing_file(label: str, value: str | None) -> str:
    if not value:
        raise SystemExit(f"[ERROR] missing required {label}; no backup/convenience fallback is allowed")
    path = Path(value)
    if not path.exists() or not path.is_file():
        raise SystemExit(
            f"[ERROR] required {label} not found: {path}\n"
            f"        No backup/convenience firmware fallback is allowed for proof runs."
        )
    return str(path)


def require_proof_pack_artifacts(proof_pack: dict[str, Any], firmware_image: str | None = None) -> tuple[str, str | None]:
    fw_block = proof_pack.get("firmware") or {}
    fw = firmware_image or fw_block.get("firmware_image") or fw_block.get("firmware_rel")
    asm = fw_block.get("disassembly") or fw_block.get("asm") or fw_block.get("asm_rel")
    fw = require_existing_file("proof firmware image", fw)
    if asm:
        asm = require_existing_file("proof disassembly", asm)
    else:
        raise SystemExit("[ERROR] proof pack has no firmware.disassembly; no backup/convenience assembly fallback is allowed")
    return fw, asm

def safe_name(text: Any) -> str:
    out = "".join(ch if ch.isalnum() or ch in "._-" else "_" for ch in str(text))
    return out.strip("._-") or "lba_cpu_trace"


def make_read10_cdb(lba: int, blocks: int = 1) -> str:
    cdb = bytearray(10)
    cdb[0] = 0x28
    cdb[2:6] = int(lba).to_bytes(4, "big")
    cdb[7:9] = int(blocks).to_bytes(2, "big")
    return hx(cdb)


def pattern_payload(lba: int, blocks: int = 1, sector_size: int = 2048) -> str:
    # 2048-byte payload is enough for READ(10) data-in completion while keeping
    # content obviously synthetic if it leaks into traces.
    total = max(1, blocks) * sector_size
    data = bytes(((i + lba) & 0xFF) for i in range(total))
    return hx(data)


def proof_lba_finding(proof_pack: dict[str, Any]) -> dict[str, Any]:
    for f in proof_pack.get("findings") or proof_pack.get("proofs") or []:
        dyn = f.get("dynamic_expectation") or f.get("dynamic_test") or {}
        if "lba" in str(f.get("type") or "").lower() or str(dyn.get("run_mode") or "").lower() == "offline_cpu_trace":
            return dict(f)
    raise SystemExit("proof pack contains no firmware LBA/offline_cpu_trace finding")


def notable_pcs(finding: dict[str, Any]) -> dict[str, list[str]]:
    raw = (((finding.get("static_evidence") or {}).get("summary") or {}).get("notable_pcs") or {})
    out: dict[str, list[str]] = {k: list(v) for k, v in raw.items() if isinstance(v, list)}
    defaults = {
        "read_path": [
            "0x90020152", "0x90020156", "0x9002015D", "0x90020166", "0x9002016C",
            "0x90020173", "0x9002017B", "0x90020181", "0x90020189", "0x90020199",
            "0x9001E3E0", "0x9001E4E8", "0x9001E541", "0x9001EEE9",
            "0x900244AD", "0x900244B7", "0x900246DB", "0x9002478A",
        ]
    }
    for k, vals in defaults.items():
        out.setdefault(k, [])
        for v in vals:
            if v not in out[k]:
                out[k].append(v)
    return out


def breakpoints_for(finding: dict[str, Any]) -> list[dict[str, Any]]:
    caps = [
        {"name": "dfa0_cdb", "address": "0x0000DFA0", "length": 12},
        {"name": "shadow_05d0", "address": "0x000005D0", "length": 12},
        {"name": "assembled_05ec", "address": "0x000005EC", "length": 4},
        {"name": "packet_05f0", "address": "0x000005F0", "length": 16},
        {"name": "phase_06f1", "address": "0x000006F1", "length": 16},
        {"name": "lba_state_0760", "address": "0x00000760", "length": 8},
    ]
    seen: set[int] = set()
    bps: list[dict[str, Any]] = []
    for kind, pcs in notable_pcs(finding).items():
        for p in pcs:
            pc = parse_int(p)
            if pc in seen:
                continue
            seen.add(pc)
            bps.append({
                "address": f"0x{pc:08X}",
                "name": f"lba_{kind}_0x{pc:08X}",
                "capture_memory": caps,
            })
    return bps


def watchpoints() -> list[dict[str, Any]]:
    return [
        {"name": "watch_dfa0_cdb", "start": "0x0000DFA0", "size": "0x10", "kind": "access", "description": "firmware-visible packet/CDB bytes"},
        {"name": "watch_df30_packet_frame", "start": "0x0000DF30", "size": "0x10", "kind": "access", "description": "packet frame/tail source"},
        {"name": "watch_05d0_cdb_shadow", "start": "0x000005D0", "size": "0x10", "kind": "access", "description": "firmware shadow of CDB bytes"},
        {"name": "watch_05ec_assembled_lba", "start": "0x000005EC", "size": "0x04", "kind": "access", "description": "assembled LBA candidate"},
        {"name": "watch_05f0_packet_state", "start": "0x000005F0", "size": "0x20", "kind": "access", "description": "packet/transfer state candidate"},
        {"name": "watch_06f1_packet_phase", "start": "0x000006F1", "size": "0x20", "kind": "access", "description": "packet phase/DMA state candidate"},
        {"name": "watch_0760_lba_state", "start": "0x00000760", "size": "0x08", "kind": "access", "description": "current/cached LBA state candidates"},
        {"name": "watch_d904_d905_status", "start": "0x0000D904", "size": "0x02", "kind": "access", "description": "packet status"},
        {"name": "watch_d950_host_status", "start": "0x0000D950", "size": "0x01", "kind": "access", "description": "host packet status"},
    ]


def probe_reads(finding: dict[str, Any]) -> list[dict[str, Any]]:
    dyn = finding.get("dynamic_expectation") or finding.get("dynamic_test") or {}
    probes = dyn.get("probe_reads") or []
    out = []
    for p in probes:
        if isinstance(p, dict):
            out.append({"host_lba": parse_int(p.get("host_lba", p.get("lba", 0))), "blocks": parse_int(p.get("blocks", 1), 1)})
        else:
            out.append({"host_lba": parse_int(p), "blocks": 1})
    return out or [{"host_lba": x, "blocks": 1} for x in (0, 32, 33, 306111, 306112, 1913889)]


def trim_to_unlock_commands(commands: list[dict[str, Any]]) -> list[dict[str, Any]]:
    # Keep through sticky descramble/unlocked capacity if present.  This prevents
    # the CPU trace from running the entire Windows probe matrix before the LBA probes.
    keep = []
    for c in commands:
        keep.append(deepcopy(c))
        tag = str(c.get("tag") or c.get("title") or "").lower()
        if tag.startswith("step10") or "unlocked_read_capacity" in tag:
            break
    return keep or deepcopy(commands)


def build_config(base_config: dict[str, Any], proof_pack: dict[str, Any], stem: str, trace_csv: Path, *, max_cycles: int, keep_all_base_commands: bool, firmware_image: str | None = None) -> dict[str, Any]:
    finding = proof_lba_finding(proof_pack)
    cfg = deepcopy(base_config)
    cfg["test_name"] = f"{stem} LBA CPU/MMIO trace proof"
    cfg["max_cycles"] = max_cycles
    cfg["allow_version_mismatch"] = True

    # v177: proof runs must execute the exact firmware/disassembly pair named
    # by the proof pack or an explicit --firmware-image.  Do not inherit
    # convenience/backup custom_bin values from the base config.
    chosen_fw, chosen_asm = require_proof_pack_artifacts(proof_pack, firmware_image)
    cfg["custom_bin"] = chosen_fw
    cfg["proof_firmware_image"] = chosen_fw
    cfg["proof_disassembly"] = chosen_asm
    cfg["proof_binary_source"] = "proof_pack" if not firmware_image else "cli"

    trace = dict(cfg.get("trace") or {})
    trace["enabled"] = True
    trace["console"] = False
    trace["out_csv"] = str(trace_csv)
    trace["breakpoints"] = breakpoints_for(finding)
    trace["watchpoints"] = watchpoints()
    cfg["trace"] = trace

    model = dict(cfg.get("gdr_packet_session_model") or trace.get("gdr_packet_session_model") or {})
    commands = list(model.get("commands") or [])
    commands = deepcopy(commands if keep_all_base_commands else trim_to_unlock_commands(commands))
    for p in probe_reads(finding):
        lba = int(p["host_lba"])
        blocks = int(p.get("blocks") or 1)
        commands.append({
            "tag": f"lba_trace_read10_lba{lba}_blocks{blocks}",
            "title": f"LBA trace READ10 host_lba={lba} blocks={blocks}",
            "direction": "in",
            "cdb_hex": make_read10_cdb(lba, blocks),
            "data_in_hex": pattern_payload(lba, blocks),
        })
    model["commands"] = commands
    model["enabled"] = True
    model["emit_events"] = True
    model["stop_on_complete"] = True
    cfg["gdr_packet_session_model"] = model
    # Keep mirrored trace-local block for older emulator versions.
    trace["gdr_packet_session_model"] = model
    return cfg


def run_cmd(cmd: list[str], cwd: Path | None, log_path: Path) -> int:
    print("[RUN]", " ".join(str(x) for x in cmd))
    proc = subprocess.run(cmd, cwd=str(cwd) if cwd else None, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    log_path.parent.mkdir(parents=True, exist_ok=True)
    log_path.write_text(proc.stdout, encoding="utf-8")
    sys.stdout.write(proc.stdout)
    return int(proc.returncode)


def write_summary_md(path: Path, *, stem: str, config_path: Path, trace_csv: Path, proof_json: Path | None, proof_md: Path | None, run_rc: int | None) -> None:
    lines = ["# HLDS v170 LBA CPU/MMIO trace proof", ""]
    lines.append(f"stem={stem}")
    lines.append(f"config={config_path}")
    lines.append(f"trace_csv={trace_csv}")
    if run_rc is not None:
        lines.append(f"emulator_returncode={run_rc}")
    if proof_json:
        lines.append(f"proof_json={proof_json}")
    if proof_md:
        lines.append(f"proof_md={proof_md}")
    lines.append("")
    lines.append("This is an offline firmware evidence run. It does not start iSCSI and does not use Windows. It is intended to prove whether the firmware path actually parses READ CDB LBA bytes and updates/submits LBA/window state, rather than relying on the high-level bridge sidecar mapping.")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    ap = argparse.ArgumentParser(description="Build/run GDR-8050L offline LBA CPU/MMIO trace proof configs")
    ap.add_argument("--base-config", required=True, help="known-good persistent GDR config, e.g. gdr8050l_auth_v166_windows_iscsi_config.json")
    ap.add_argument("--proof-pack", required=True, help="v169 LBA math candidates proof pack")
    ap.add_argument("--stem", default="gdr8050l_0l23_v170_lba_cpu_trace")
    ap.add_argument("--out-dir", default=None)
    ap.add_argument("--run", action="store_true", help="run mn103s_emu.py after generating the config")
    ap.add_argument("--validate", action="store_true", help="validate the generated trace after run, or an existing --trace-csv")
    ap.add_argument("--trace-csv", default=None, help="existing trace CSV to validate instead of generated trace path")
    ap.add_argument("--python", default=sys.executable)
    ap.add_argument("--emulator", default="mn103s_emu.py")
    ap.add_argument("--cwd", default=None)
    ap.add_argument("--max-cycles", type=int, default=360000)
    ap.add_argument("--firmware-image", default=None, help="explicit proof firmware image; default is proof_pack.firmware.firmware_image; missing files are fatal")
    ap.add_argument("--keep-all-base-commands", action="store_true", help="do not trim base command matrix to auth/unlock commands before LBA probes")
    args = ap.parse_args()

    base_config_path = Path(args.base_config)
    proof_pack_path = Path(args.proof_pack)
    base_config = load_json_or_jsonc(base_config_path)
    proof_pack = load_json_or_jsonc(proof_pack_path)
    stem = safe_name(args.stem)
    out_dir = Path(args.out_dir or (Path("emutests") / stem))
    out_dir.mkdir(parents=True, exist_ok=True)
    config_path = out_dir / f"{stem}.config.json"
    trace_csv = Path(args.trace_csv) if args.trace_csv else out_dir / f"{stem}.trace.csv"
    # Use an absolute trace output path in the emulator config.  mn103s_emu.py
    # resolves some relative paths against the active config directory, while
    # the builder/validator used the caller CWD.  Absolute paths keep the
    # writer and validator on the same file.
    trace_csv_for_config = trace_csv.resolve()

    cfg = build_config(
        base_config, proof_pack, stem, trace_csv_for_config,
        max_cycles=args.max_cycles,
        keep_all_base_commands=args.keep_all_base_commands,
        firmware_image=args.firmware_image,
    )
    trace_csv = trace_csv_for_config
    write_json(config_path, cfg)
    print(f"[OK] wrote LBA CPU trace config: {config_path}")
    print(f"[INFO] proof firmware image: {cfg.get('custom_bin', '<none>')}")
    print(f"[INFO] trace CSV: {trace_csv}")

    run_rc: int | None = None
    if args.run:
        run_rc = run_cmd([args.python, args.emulator, str(config_path)], Path(args.cwd) if args.cwd else None, out_dir / f"{stem}.run.log")

    proof_json = proof_md = None
    if args.validate or args.run:
        if not trace_csv.exists():
            print(f"[WARN] trace CSV does not exist yet; skipping validation: {trace_csv}")
        else:
            report = validate_lba_trace(trace_csv, proof_pack)
            proof_json = out_dir / f"{stem}.lba_trace_proof.json"
            proof_md = out_dir / f"{stem}.lba_trace_proof.md"
            write_json(proof_json, report)
            write_lba_md(proof_md, report)
            print(f"[OK] wrote LBA CPU trace proof JSON: {proof_json}")
            print(f"[OK] wrote LBA CPU trace proof MD: {proof_md}")

    write_summary_md(out_dir / "README_v170_run.md", stem=stem, config_path=config_path, trace_csv=trace_csv, proof_json=proof_json, proof_md=proof_md, run_rc=run_rc)
    print(f"[OK] wrote summary: {out_dir / 'README_v170_run.md'}")
    return 0 if run_rc in (None, 0) else run_rc


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