#!/usr/bin/env python3
"""HLDS/GDR-8050L direct firmware LBA function probes (v172).

This is intentionally *not* a sidecar/bridge proof.  It starts small firmware
routines directly with the firmware-visible CDB bytes seeded at DFA0..DFAB and
observes the firmware code that:

  * copies READ CDB LBA bytes into 0x05EC / packet state, and
  * computes READ LBA + 150 in the 0x90020152 helper.

These probes do not claim that the full host-interface route has reached these
routines yet.  They prove what the candidate routines do when executed.
"""
from __future__ import annotations

import argparse
import csv
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

VERSION = 1

READ10_PACKET_INGRESS_PC = 0x9001E4E5
READ_LBA_PLUS150_PC = 0x90020152
ROM_BASE = "0x90000000"



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 parse_int(x: Any, default: int = 0) -> int:
    if x is None:
        return default
    if isinstance(x, int):
        return x
    s = str(x).strip()
    if not s:
        return default
    return int(s, 0)


def load_json(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8", errors="ignore") as f:
        return json.load(f)


def write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="\n") as f:
        json.dump(obj, f, indent=2)
        f.write("\n")


def cdb_read10(host_lba: int, blocks: int = 1) -> List[int]:
    host_lba &= 0xFFFFFFFF
    blocks &= 0xFFFF
    return [
        0x28,
        0x00,
        (host_lba >> 24) & 0xFF,
        (host_lba >> 16) & 0xFF,
        (host_lba >> 8) & 0xFF,
        host_lba & 0xFF,
        0x00,
        (blocks >> 8) & 0xFF,
        blocks & 0xFF,
        0x00,
        0x00,
        0x00,
    ]


def hex_bytes(xs: List[int]) -> str:
    return " ".join(f"{int(x) & 0xFF:02X}" for x in xs)


def extract_probe_reads(proof_pack: Dict[str, Any]) -> List[Dict[str, int]]:
    reads: List[Dict[str, int]] = []
    for finding in proof_pack.get("findings", []) or []:
        dyn = finding.get("dynamic_expectation", {}) or {}
        for r in dyn.get("probe_reads", []) or []:
            try:
                host = parse_int(r.get("host_lba"))
            except Exception:
                continue
            bridge = r.get("expected_current_bridge_iso_lba")
            try:
                bridge_i = parse_int(bridge) if bridge is not None else host
            except Exception:
                bridge_i = host
            if host not in [x["host_lba"] for x in reads]:
                reads.append({"host_lba": host, "expected_current_bridge_iso_lba": bridge_i})
    if not reads:
        for host in [0, 32, 33, 306111, 306112, 1913889]:
            reads.append({"host_lba": host, "expected_current_bridge_iso_lba": host + 306080})
    return reads


def base_config(
    *,
    test_name: str,
    firmware_image: str,
    entry_point: int,
    max_cycles: int,
    out_csv: Path,
    sp: int = 0x003FF000,
) -> Dict[str, Any]:
    return {
        "test_name": test_name,
        "workspace": str(Path.cwd()),
        "target_profile": "GDR-8050L",
        "custom_bin": firmware_image,
        "custom_base_addr": ROM_BASE,
        "entry_point": f"0x{entry_point:08X}",
        "header_skip_bytes": "0x0",
        "config_version": "v100",
        "requires_emulator_version": "V7-direct-config-trace-v100-gcc-version-sync",
        "allow_version_mismatch": True,
        "max_cycles": max_cycles,
        "initial_registers": {"SP": f"0x{sp:08X}"},
        "memory_aliases": [
            {
                "name": "SRAM alias 0x80000000->0x00000000",
                "from": "0x80000000",
                "to": "0x00000000",
                "size": "0x400000",
            }
        ],
        "trace": {
            "enabled": True,
            "live_instruction_log": False,
            "out_csv": str(out_csv.resolve()),
            "watchpoints": [
                {"name": "watch_dfa0_cdb", "start": "0x0000DFA0", "size": "0x10", "kind": "access", "description": "firmware-visible CDB bytes"},
                {"name": "watch_05d0_shadow", "start": "0x000005D0", "size": "0x10", "kind": "access", "description": "firmware CDB shadow"},
                {"name": "watch_05ec_lba_state", "start": "0x000005EC", "size": "0x20", "kind": "access", "description": "assembled LBA / packet state"},
                {"name": "watch_06f1_phase", "start": "0x000006F1", "size": "0x10", "kind": "access", "description": "packet/DMA phase"},
            ],
        },
    }


def config_packet_ingress(host_lba: int, firmware_image: str, out_dir: Path, stem: str) -> Tuple[Path, Path, Dict[str, Any]]:
    name = f"{stem}.ingress_lba{host_lba}"
    trace_csv = out_dir / f"{name}.trace.csv"
    cfg = base_config(test_name=f"{name} direct READ packet ingress probe", firmware_image=firmware_image,
                      entry_point=READ10_PACKET_INGRESS_PC, max_cycles=160, out_csv=trace_csv)
    cfg["memory_injections"] = [
        {"address": "0x0000DFA0", "description": f"READ10 CDB host_lba={host_lba}", "data": cdb_read10(host_lba, 1)}
    ]
    cfg["trace"]["breakpoints"] = [
        {"address": "0x9001E4E8", "name": "ingress_read10_cmp"},
        {"address": "0x9001E527", "name": "ingress_cdb_lba_read"},
        {"address": "0x9001E541", "name": "ingress_write_05ec"},
        {"address": "0x9001E546", "name": "ingress_write_05f0"},
        {"address": "0x9001E551", "name": "ingress_shadow_opcode"},
        {"address": "0x9001E556", "name": "ingress_write_06f1"},
    ]
    cfg_path = out_dir / f"{name}.config.json"
    return cfg_path, trace_csv, cfg


def config_plus150(host_lba: int, compare_value: int, mode: str, firmware_image: str, out_dir: Path, stem: str) -> Tuple[Path, Path, Dict[str, Any]]:
    name = f"{stem}.plus150_lba{host_lba}_{mode}"
    trace_csv = out_dir / f"{name}.trace.csv"
    cfg = base_config(test_name=f"{name} direct READ LBA +150 helper probe", firmware_image=firmware_image,
                      entry_point=READ_LBA_PLUS150_PC, max_cycles=160, out_csv=trace_csv)
    cfg["initial_registers"]["D0"] = f"0x{compare_value & 0xFFFFFFFF:08X}"
    cfg["memory_injections"] = [
        {"address": "0x0000DFA0", "description": f"READ10 CDB host_lba={host_lba}; compare={compare_value}", "data": cdb_read10(host_lba, 1)}
    ]
    cfg["trace"]["breakpoints"] = [
        {"address": "0x90020156", "name": "plus150_read10_cmp"},
        {"address": "0x9002016C", "name": "plus150_cdb_lba_read"},
        {"address": "0x90020189", "name": "plus150_add_150"},
        {"address": "0x900201BA", "name": "plus150_compare"},
        {"address": "0x900201BF", "name": "plus150_match_return_1"},
        {"address": "0x900201C3", "name": "plus150_mismatch_return_0"},
    ]
    cfg_path = out_dir / f"{name}.config.json"
    return cfg_path, trace_csv, cfg


def read_trace_rows(path: Path) -> List[Dict[str, str]]:
    if not path.exists():
        return []
    with path.open("r", encoding="utf-8", errors="ignore", newline="") as f:
        return list(csv.DictReader(f))


def as_int_field(row: Dict[str, str], key: str, default: Optional[int] = None) -> Optional[int]:
    v = row.get(key)
    if v is None or v == "":
        return default
    try:
        return int(v, 0)
    except Exception:
        return default


def little32_from_byte_writes(rows: List[Dict[str, str]], base: int, pc: Optional[int] = None) -> Optional[int]:
    b = {}
    for r in rows:
        if r.get("event") != "watch_write":
            continue
        addr = as_int_field(r, "addr")
        if addr is None or not (base <= addr < base + 4):
            continue
        if pc is not None and as_int_field(r, "pc") != pc:
            continue
        val = as_int_field(r, "value")
        if val is not None:
            b[addr - base] = val & 0xFF
    if not b:
        return None
    return sum((b.get(i, 0) & 0xFF) << (8 * i) for i in range(4))


def validate_ingress(rows: List[Dict[str, str]], host_lba: int) -> Dict[str, Any]:
    pcs = {as_int_field(r, "pc") for r in rows if r.get("event") == "break"}
    lba05ec = little32_from_byte_writes(rows, 0x05EC, 0x9001E541)
    wrote_05f0_0e = any(r.get("event") == "watch_write" and as_int_field(r, "addr") == 0x05F0 and as_int_field(r, "value") == 0x0E for r in rows)
    wrote_05d0_28 = any(r.get("event") == "watch_write" and as_int_field(r, "addr") == 0x05D0 and as_int_field(r, "value") == 0x28 for r in rows)
    wrote_06f1_04 = any(r.get("event") == "watch_write" and as_int_field(r, "addr") == 0x06F1 and as_int_field(r, "value") == 0x04 for r in rows)
    cdb_reads = [r for r in rows if r.get("event") == "watch_read" and 0xDFA2 <= (as_int_field(r, "addr", -1) or -1) <= 0xDFA5]
    proven = (0x9001E541 in pcs and lba05ec == host_lba and wrote_05f0_0e and wrote_05d0_28 and wrote_06f1_04 and len(cdb_reads) >= 4)
    return {
        "outcome": "PROVEN" if proven else "CONTRADICTED",
        "break_pcs": [f"0x{x:08X}" for x in sorted([p for p in pcs if p is not None])],
        "cdb_lba_byte_reads": len(cdb_reads),
        "lba_05ec": lba05ec,
        "wrote_05f0_0e": wrote_05f0_0e,
        "wrote_05d0_28": wrote_05d0_28,
        "wrote_06f1_04": wrote_06f1_04,
    }


def validate_plus150(rows: List[Dict[str, str]], host_lba: int, compare_value: int, expect_match: bool) -> Dict[str, Any]:
    pcs = {as_int_field(r, "pc") for r in rows if r.get("event") == "break"}
    d0_at_add = None
    d0_after_add = None
    d1_at_compare = None
    final_d0 = None
    for r in rows:
        pc = as_int_field(r, "pc")
        if r.get("event") == "break" and pc == 0x90020189:
            d0_at_add = as_int_field(r, "d0")
        if r.get("event") == "break" and pc == 0x900201BA:
            d0_after_add = as_int_field(r, "d0")
        final_d0 = as_int_field(r, "d0", final_d0)
        if pc == 0x900201BA:
            d1_at_compare = as_int_field(r, "d1")
    hit_core = 0x90020156 in pcs and 0x9002016C in pcs and 0x90020189 in pcs and 0x900201BA in pcs
    computed_ok = d0_after_add == (host_lba + 150) & 0xFFFFFFFF
    match_hit = 0x900201BF in pcs
    mismatch_hit = 0x900201C3 in pcs
    path_ok = (match_hit and expect_match) or (mismatch_hit and not expect_match)
    # The trace breakpoint fires before the final mov/clr has executed, so the
    # path hit is the reliable proof.  final_d0 is included only as context.
    proven = hit_core and computed_ok and path_ok
    return {
        "outcome": "PROVEN" if proven else "CONTRADICTED",
        "break_pcs": [f"0x{x:08X}" for x in sorted([p for p in pcs if p is not None])],
        "host_lba": host_lba,
        "compare_value": compare_value,
        "expected_match": expect_match,
        "d0_at_add": d0_at_add,
        "computed_lba_plus_150": d0_after_add,
        "expected_lba_plus_150": (host_lba + 150) & 0xFFFFFFFF,
        "final_d0": final_d0,
        "match_path_hit": match_hit,
        "mismatch_path_hit": mismatch_hit,
    }


def run_config(python_exe: str, emulator: Path, cfg_path: Path, run_log: Path, cwd: Path) -> int:
    cmd = [python_exe, str(emulator), str(cfg_path)]
    run_log.parent.mkdir(parents=True, exist_ok=True)
    with run_log.open("w", encoding="utf-8", errors="ignore", newline="\n") as f:
        p = subprocess.run(cmd, cwd=str(cwd), stdout=f, stderr=subprocess.STDOUT)
    return p.returncode


def main(argv: Optional[List[str]] = None) -> int:
    ap = argparse.ArgumentParser(description="Run direct offline firmware LBA function probes")
    ap.add_argument("--proof-pack", required=True)
    ap.add_argument("--stem", default="gdr8050l_0l23_v172_direct_lba_function_probe")
    ap.add_argument("--out-dir", default=None)
    ap.add_argument("--firmware-image", default=None)
    ap.add_argument("--emulator", default="mn103s_emu.py")
    ap.add_argument("--python", default=sys.executable)
    ap.add_argument("--cwd", default=".")
    ap.add_argument("--run", action="store_true")
    ap.add_argument("--validate", action="store_true")
    ap.add_argument("--host-lba", action="append", default=[], help="Limit/add host LBA probe; may be decimal or hex")
    args = ap.parse_args(argv)

    proof_path = Path(args.proof_pack)
    proof_pack = load_json(proof_path)
    fw_block = proof_pack.get("firmware", {}) or {}
    firmware_image = require_existing_file("proof firmware image", args.firmware_image or fw_block.get("firmware_image"))
    require_existing_file("proof disassembly", fw_block.get("disassembly"))

    reads = extract_probe_reads(proof_pack)
    if args.host_lba:
        wanted = {parse_int(x) for x in args.host_lba}
        existing = {r["host_lba"] for r in reads}
        for host in wanted - existing:
            reads.append({"host_lba": host, "expected_current_bridge_iso_lba": host + 306080})
        reads = [r for r in reads if r["host_lba"] in wanted]

    cwd = Path(args.cwd).resolve()
    emulator = Path(args.emulator)
    if not emulator.is_absolute():
        emulator = cwd / emulator
    out_dir = Path(args.out_dir) if args.out_dir else Path("emutests") / args.stem
    out_dir.mkdir(parents=True, exist_ok=True)

    manifest: Dict[str, Any] = {
        "format": "hlds-lba-direct-function-probe-report",
        "version": VERSION,
        "stem": args.stem,
        "proof_pack": str(proof_path),
        "firmware_image": str(firmware_image),
        "probes": [],
        "summary": {},
    }

    tasks: List[Tuple[str, int, str, bool, Path, Path, Dict[str, Any], str]] = []
    # Ingress probe for each host LBA.
    for r in reads:
        host = int(r["host_lba"])
        cfg_path, trace_csv, cfg = config_packet_ingress(host, str(firmware_image), out_dir, args.stem)
        write_json(cfg_path, cfg)
        tasks.append(("packet_ingress", host, "ingress", True, cfg_path, trace_csv, cfg, str(out_dir / (cfg_path.stem + ".run.log"))))
        # LBA+150 helper: positive and two negative controls.
        modes = [
            ("host_plus_150", host + 150, True),
            ("host_identity", host, False),
            ("bridge_iso", int(r.get("expected_current_bridge_iso_lba", host + 306080)), False),
        ]
        for mode, compare, expect in modes:
            cfg_path, trace_csv, cfg = config_plus150(host, compare, mode, str(firmware_image), out_dir, args.stem)
            write_json(cfg_path, cfg)
            tasks.append(("plus150", host, mode, expect, cfg_path, trace_csv, cfg, str(out_dir / (cfg_path.stem + ".run.log"))))

    for kind, host, mode, expect, cfg_path, trace_csv, cfg, run_log_s in tasks:
        rc: Optional[int] = None
        if args.run:
            print(f"[RUN] {cfg_path}")
            rc = run_config(args.python, emulator, cfg_path, Path(run_log_s), cwd)
        result: Dict[str, Any]
        if args.validate:
            rows = read_trace_rows(trace_csv)
            if kind == "packet_ingress":
                result = validate_ingress(rows, host)
            else:
                compare = parse_int(cfg["initial_registers"]["D0"])
                result = validate_plus150(rows, host, compare, expect)
        else:
            result = {"outcome": "GENERATED"}
        result.update({
            "kind": kind,
            "host_lba": host,
            "mode": mode,
            "config": str(cfg_path),
            "trace_csv": str(trace_csv),
            "run_log": run_log_s,
            "returncode": rc,
        })
        manifest["probes"].append(result)
        print(f"[PROBE] {kind} host_lba={host} mode={mode} outcome={result.get('outcome')}")

    total = len(manifest["probes"])
    proven = sum(1 for p in manifest["probes"] if p.get("outcome") == "PROVEN")
    contradicted = sum(1 for p in manifest["probes"] if p.get("outcome") == "CONTRADICTED")
    manifest["summary"] = {"total": total, "proven": proven, "contradicted": contradicted}

    report_json = out_dir / f"{args.stem}.direct_lba_function_probe_report.json"
    write_json(report_json, manifest)

    report_md = out_dir / f"{args.stem}.direct_lba_function_probe_report.md"
    with report_md.open("w", encoding="utf-8", newline="\n") as f:
        f.write("# HLDS v172 direct firmware LBA function probe report\n\n")
        f.write("This is an offline direct-function proof. It does not prove the full host packet route; it proves what the candidate firmware routines do when entered with firmware-visible DFA0 CDB bytes seeded.\n\n")
        f.write(f"stem={args.stem}\n\n")
        f.write(f"summary: proven={proven}/{total}, contradicted={contradicted}\n\n")
        f.write("| kind | host_lba | mode | outcome | key result |\n")
        f.write("|---|---:|---|---|---|\n")
        for p in manifest["probes"]:
            if p["kind"] == "packet_ingress":
                key = f"05EC={p.get('lba_05ec')} 05F0=0x0E:{p.get('wrote_05f0_0e')} 06F1=0x04:{p.get('wrote_06f1_04')}"
            else:
                key = f"compare={p.get('compare_value')} computed+150={p.get('computed_lba_plus_150')} final_d0={p.get('final_d0')}"
            f.write(f"| `{p['kind']}` | {p['host_lba']} | `{p['mode']}` | {p.get('outcome')} | {key} |\n")
        f.write("\nInterpretation: `packet_ingress` proves the READ CDB bytes at DFA0 are assembled into 0x05EC and packet state is written. `plus150` proves the 0x90020152 helper compares an input value against CDB LBA + 150; the bridge ISO/window value is a negative control and should not match.\n")

    print(f"[OK] wrote report JSON: {report_json}")
    print(f"[OK] wrote report MD: {report_md}")
    return 0 if contradicted == 0 else 2


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