#!/usr/bin/env python3
"""Validate firmware-side LBA proof evidence from an MN103S trace CSV.

This is deliberately stricter than the high-level bridge read-map checks.  A
bridge/sidecar response can show that a READ works; it does not prove that the
GDR-8050L firmware parsed the CDB LBA and submitted a translated/source LBA to
its packet/DMA hardware path.  This validator consumes the CPU/MMIO trace emitted
by mn103s_emu.py and maps it back to analyzer proof-pack findings.
"""
from __future__ import annotations

import argparse
import csv
import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Iterable

try:
    from hlds_profiles import load_json_or_jsonc, write_json
except Exception:  # standalone fallback
    def _strip_jsonc(text: str) -> str:
        text = re.sub(r"/\*.*?\*/", "", text, flags=re.S)
        out = []
        for line in text.splitlines():
            s = line
            if "//" in s:
                s = s.split("//", 1)[0]
            out.append(s)
        return "\n".join(out)

    def load_json_or_jsonc(path: str | Path) -> dict[str, Any]:
        return json.loads(_strip_jsonc(Path(path).read_text(encoding="utf-8")))

    def write_json(path: str | Path, data: Any) -> None:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        Path(path).write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def parse_int(v: Any, default: int | None = None) -> int:
    if v is None or v == "":
        if default is None:
            raise ValueError("empty integer")
        return default
    if isinstance(v, int):
        return v
    s = str(v).strip()
    if not s:
        if default is None:
            raise ValueError("empty integer")
        return default
    return int(s, 0)


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


def parse_hex_bytes(text: Any) -> bytes:
    if text is None:
        return b""
    s = str(text).strip().replace("0x", "").replace("0X", "")
    s = "".join(ch for ch in s if ch in "0123456789abcdefABCDEF")
    if len(s) % 2:
        s = "0" + s
    return bytes.fromhex(s) if s else b""


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_trace"


def normalize_addr(v: Any) -> int | None:
    try:
        return parse_int(v)
    except Exception:
        return None


def overlap(row_addr: int | None, row_size: int, start: int, end: int) -> bool:
    if row_addr is None:
        return False
    row_end = row_addr + max(1, row_size) - 1
    return not (row_end < start or row_addr > end)


@dataclass
class CommandWindow:
    index: int
    name: str
    start_row: int
    end_row: int | None = None
    cdb: bytes = b""
    rows: list[dict[str, str]] = field(default_factory=list)

    @property
    def lba(self) -> int | None:
        if len(self.cdb) >= 10 and self.cdb[0] == 0x28:
            return int.from_bytes(self.cdb[2:6], "big")
        if len(self.cdb) >= 12 and self.cdb[0] == 0xA8:
            return int.from_bytes(self.cdb[2:6], "big")
        return None

    @property
    def opcode(self) -> int | None:
        return self.cdb[0] if self.cdb else None


def read_trace_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8", errors="ignore") as f:
        return [dict(r) for r in csv.DictReader(f)]


def command_windows(rows: list[dict[str, str]]) -> list[CommandWindow]:
    windows: list[CommandWindow] = []
    current: CommandWindow | None = None
    for i, row in enumerate(rows):
        event = str(row.get("event") or "")
        if event == "gdr_session_command_arm":
            if current is not None:
                current.end_row = i - 1
                windows.append(current)
            current = CommandWindow(
                index=len(windows),
                name=str(row.get("name") or f"cmd{len(windows):02d}"),
                start_row=i,
                cdb=parse_hex_bytes(row.get("bytes_hex") or ""),
                rows=[row],
            )
            continue
        if current is not None:
            current.rows.append(row)
            if event == "gdr_session_command_complete":
                current.end_row = i
                windows.append(current)
                current = None
    if current is not None:
        current.end_row = len(rows) - 1
        windows.append(current)
    if not windows:
        # Non-session trace fallback: one synthetic window over all rows.
        windows.append(CommandWindow(index=0, name="trace", start_row=0, end_row=len(rows)-1, rows=rows))
    return windows


def lba_findings(proof_pack: dict[str, Any]) -> list[dict[str, Any]]:
    out = []
    for f in proof_pack.get("findings") or proof_pack.get("proofs") or []:
        dyn = f.get("dynamic_expectation") or f.get("dynamic_test") or {}
        ftype = str(f.get("type") or "").lower()
        if "lba" in ftype or str(dyn.get("run_mode") or "").lower() == "offline_cpu_trace":
            out.append(dict(f))
    return out


def notable_pcs(finding: dict[str, Any]) -> dict[str, set[int]]:
    se = finding.get("static_evidence") or {}
    summary = se.get("summary") or {}
    raw = summary.get("notable_pcs") or {}
    groups: dict[str, set[int]] = {}
    for k, vals in raw.items():
        groups[k] = {parse_int(v) for v in vals if str(v).strip()}
    # Hard defaults from v169 locator in case older proof packs are minimal.
    groups.setdefault("read10_opcode_compare_0x28", set()).update({0x9001E3E0, 0x9001E4E8, 0x90020156})
    groups.setdefault("read12_opcode_compare_0xA8", set()).update({0x9001E3E7, 0x9001E500, 0x9002015D})
    groups.setdefault("read_cd_or_related_opcode_compare_0xBE", set()).update({0x90020166})
    groups.setdefault("candidate_addr_dfa2_packet_cdb_lba_b0_or_shadow", set()).update({0x9002016C})
    groups.setdefault("candidate_addr_dfa3_packet_cdb_lba_b1_or_shadow", set()).update({0x90020173})
    groups.setdefault("candidate_addr_dfa4_packet_cdb_lba_b2_or_shadow", set()).update({0x9002017B})
    groups.setdefault("candidate_addr_dfa5_packet_cdb_lba_b3_or_shadow", set()).update({0x90020181})
    groups.setdefault("candidate_addr_05ec_assembled_lba_candidate", set()).update({0x9001E541, 0x9001EEE9, 0x900244AD, 0x900244B7})
    groups.setdefault("candidate_read_lba_plus_150", set()).update({0x90020189, 0x900246DB, 0x9002478A})
    return groups


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


def pcs_in_rows(rows: list[dict[str, str]], pcset: set[int]) -> list[str]:
    hits = []
    for r in rows:
        pc = normalize_addr(r.get("pc"))
        if pc in pcset:
            hits.append(f"0x{pc:08X}")
    return sorted(set(hits))


def count_watch(rows: list[dict[str, str]], event_prefix: str, start: int, end: int) -> tuple[int, list[dict[str, Any]]]:
    hits: list[dict[str, Any]] = []
    for r in rows:
        if not str(r.get("event") or "").startswith(event_prefix):
            continue
        addr = normalize_addr(r.get("addr"))
        phys = normalize_addr(r.get("addr_physical"))
        try:
            size = int(str(r.get("size") or "1"), 0)
        except Exception:
            size = 1
        if overlap(addr, size, start, end) or overlap(phys, size, start, end):
            hits.append({
                "pc": r.get("pc"),
                "addr": r.get("addr"),
                "addr_physical": r.get("addr_physical"),
                "size": r.get("size"),
                "value": r.get("value"),
                "bytes_hex": r.get("bytes_hex"),
                "name": r.get("name"),
                "event": r.get("event"),
            })
    return len(hits), hits[:24]


def count_events(rows: list[dict[str, str]], event_name: str) -> tuple[int, list[dict[str, Any]]]:
    hits: list[dict[str, Any]] = []
    for r in rows:
        if str(r.get("event") or "") != event_name:
            continue
        hits.append({
            "pc": r.get("pc"),
            "addr": r.get("addr"),
            "addr_physical": r.get("addr_physical"),
            "size": r.get("size"),
            "value": r.get("value"),
            "bytes_hex": r.get("bytes_hex"),
            "name": r.get("name"),
            "event": r.get("event"),
            "notes": r.get("notes"),
        })
    return len(hits), hits[:24]


def has_cdb_lba_bytes(rows: list[dict[str, str]], host_lba: int) -> tuple[bool, dict[str, Any]]:
    want = host_lba.to_bytes(4, "big")
    observed: dict[int, list[int]] = {0xDFA2+i: [] for i in range(4)}
    shadow: dict[int, list[int]] = {0x05D2+i: [] for i in range(4)}
    for r in rows:
        ev = str(r.get("event") or "")
        if not ev.startswith("watch_read") and not ev.startswith("watch_write"):
            continue
        addrs = [normalize_addr(r.get("addr")), normalize_addr(r.get("addr_physical"))]
        value = normalize_addr(r.get("value"))
        if value is None:
            bs = parse_hex_bytes(r.get("bytes_hex") or "")
            value = bs[0] if bs else None
        if value is None:
            continue
        for a in addrs:
            if a in observed:
                observed[a].append(value & 0xFF)
            if a in shadow:
                shadow[a].append(value & 0xFF)
    dfa_match = all(want[i] in observed[0xDFA2+i] for i in range(4))
    shadow_match = all(want[i] in shadow[0x05D2+i] for i in range(4))
    return dfa_match or shadow_match, {
        "expected_lba_bytes": hx_bytes(want),
        "dfa2_dfa5_observed": {f"0x{k:04X}": [f"0x{x:02X}" for x in vals[:16]] for k, vals in observed.items()},
        "shadow_05d2_05d5_observed": {f"0x{k:04X}": [f"0x{x:02X}" for x in vals[:16]] for k, vals in shadow.items()},
        "dfa_match": dfa_match,
        "shadow_match": shadow_match,
    }


def analyze_window(win: CommandWindow, finding: dict[str, Any]) -> dict[str, Any]:
    groups = notable_pcs(finding)
    rows = win.rows
    host_lba = win.lba
    result: dict[str, Any] = {
        "command": win.name,
        "opcode": f"0x{win.opcode:02X}" if win.opcode is not None else None,
        "cdb_hex": hx_bytes(win.cdb),
        "host_lba": host_lba,
        "row_count": len(rows),
    }
    read_handler_pcs = set().union(
        groups.get("read10_opcode_compare_0x28", set()),
        groups.get("read12_opcode_compare_0xA8", set()),
        groups.get("read_cd_or_related_opcode_compare_0xBE", set()),
    )
    result["read_handler_hits"] = pcs_in_rows(rows, read_handler_pcs)
    result["lba_byte_read_pc_hits"] = pcs_in_rows(rows, set().union(
        groups.get("candidate_addr_dfa2_packet_cdb_lba_b0_or_shadow", set()),
        groups.get("candidate_addr_dfa3_packet_cdb_lba_b1_or_shadow", set()),
        groups.get("candidate_addr_dfa4_packet_cdb_lba_b2_or_shadow", set()),
        groups.get("candidate_addr_dfa5_packet_cdb_lba_b3_or_shadow", set()),
    ))
    result["assembled_lba_pc_hits"] = pcs_in_rows(rows, groups.get("candidate_addr_05ec_assembled_lba_candidate", set()))
    result["plus_150_pc_hits"] = pcs_in_rows(rows, groups.get("candidate_read_lba_plus_150", set()))

    result["watch_dfa2_dfa5"] = count_watch(rows, "watch_read", 0xDFA2, 0xDFA5)[0]
    result["watch_05d2_05d5"] = count_watch(rows, "watch", 0x05D2, 0x05D5)[0]
    result["watch_05ec_write"] = count_watch(rows, "watch_write", 0x05EC, 0x05EF)[0]
    result["watch_05f0_write"] = count_watch(rows, "watch_write", 0x05F0, 0x05FF)[0]
    result["watch_06f1_write"] = count_watch(rows, "watch_write", 0x06F1, 0x06FF)[0]
    result["watch_0760_0767_write"] = count_watch(rows, "watch_write", 0x0760, 0x0767)[0]
    result["sample_05ec_writes"] = count_watch(rows, "watch_write", 0x05EC, 0x05EF)[1]
    result["sample_lba_state_writes"] = (count_watch(rows, "watch_write", 0x0760, 0x0767)[1] + count_watch(rows, "watch_write", 0x05F0, 0x06FF)[1])[:24]
    fw_source_count, fw_source_samples = count_events(rows, "gdr_session_firmware_lba_source")
    result["firmware_lba_source_event_count"] = fw_source_count
    result["sample_firmware_lba_source_events"] = fw_source_samples

    if host_lba is not None:
        ok, obs = has_cdb_lba_bytes(rows, host_lba)
        result["cdb_lba_bytes_observed"] = ok
        result["cdb_lba_byte_observations"] = obs
    else:
        result["cdb_lba_bytes_observed"] = False

    handler_hit = bool(result["read_handler_hits"])
    cdb_seen = bool(result["cdb_lba_bytes_observed"] or result["lba_byte_read_pc_hits"])
    lba_state = bool(result["watch_05ec_write"] or result["watch_0760_0767_write"] or result["assembled_lba_pc_hits"])
    packet_state = bool(result["watch_05f0_write"] or result["watch_06f1_write"])
    firmware_source = bool(result.get("firmware_lba_source_event_count"))
    if handler_hit and cdb_seen and lba_state and packet_state:
        result["outcome"] = "PROVEN"
        result["route_evidence_level"] = "firmware_lba_source" if firmware_source else "firmware_ingress_state"
    elif handler_hit or cdb_seen or lba_state or packet_state or firmware_source:
        result["outcome"] = "PARTIAL"
        result["route_evidence_level"] = "partial_firmware_lba_source" if firmware_source else "partial"
    else:
        result["outcome"] = "BLOCKED"
    result["checks"] = {
        "handler_hit": handler_hit,
        "cdb_lba_seen": cdb_seen,
        "lba_state_seen": lba_state,
        "packet_state_seen": packet_state,
        "firmware_lba_source_seen": firmware_source,
    }
    return result


def validate_lba_trace(trace_csv: str | Path, proof_pack: dict[str, Any]) -> dict[str, Any]:
    trace_path = Path(trace_csv)
    rows = read_trace_csv(trace_path)
    windows = command_windows(rows)
    findings = lba_findings(proof_pack)
    out: dict[str, Any] = {
        "format": "hlds-lba-cpu-trace-proof-report",
        "version": 1,
        "trace_csv": str(trace_path),
        "trace_rows": len(rows),
        "command_windows": len(windows),
        "findings": [],
    }
    for finding in findings:
        fid = str(finding.get("finding_id") or finding.get("id") or "lba_finding")
        expected_lbas = set(probe_lbas(finding))
        read_windows = [w for w in windows if w.opcode in (0x28, 0xA8)]
        if expected_lbas:
            read_windows = [w for w in read_windows if w.lba in expected_lbas]
        probes = [analyze_window(w, finding) for w in read_windows]
        missing = sorted(expected_lbas - {int(p["host_lba"]) for p in probes if p.get("host_lba") is not None})
        outcomes = {p.get("outcome") for p in probes}
        if probes and all(p.get("outcome") == "PROVEN" for p in probes) and not missing:
            outcome = "PROVEN"
        elif probes and ("PARTIAL" in outcomes or "PROVEN" in outcomes):
            outcome = "PARTIAL"
        else:
            outcome = "BLOCKED"
        out["findings"].append({
            "finding_id": fid,
            "type": finding.get("type"),
            "outcome": outcome,
            "expected_lbas": sorted(expected_lbas),
            "missing_lbas": missing,
            "probes": probes,
        })
    return out


def write_md(path: Path, report: dict[str, Any]) -> None:
    lines = ["# HLDS LBA CPU/MMIO trace proof report", ""]
    lines.append(f"trace_csv={report.get('trace_csv')}")
    lines.append(f"trace_rows={report.get('trace_rows')}")
    lines.append(f"command_windows={report.get('command_windows')}")
    lines.append("")
    for f in report.get("findings") or []:
        lines.append(f"## {f.get('finding_id')}")
        lines.append("")
        lines.append(f"outcome: **{f.get('outcome')}**")
        if f.get("missing_lbas"):
            lines.append(f"missing_lbas: {f.get('missing_lbas')}")
        lines.append("")
        lines.append("| host_lba | command | outcome | handler | cdb bytes | lba state | packet state | fw source | notable hits |")
        lines.append("|---:|---|---|---:|---:|---:|---:|---:|---|")
        for p in f.get("probes") or []:
            chk = p.get("checks") or {}
            hits = ", ".join((p.get("read_handler_hits") or []) + (p.get("assembled_lba_pc_hits") or []) + (p.get("plus_150_pc_hits") or []))
            lines.append(
                f"| {p.get('host_lba')} | `{p.get('command')}` | **{p.get('outcome')}** | "
                f"{int(bool(chk.get('handler_hit')))} | {int(bool(chk.get('cdb_lba_seen')))} | "
                f"{int(bool(chk.get('lba_state_seen')))} | {int(bool(chk.get('packet_state_seen')))} | "
                f"{int(bool(chk.get('firmware_lba_source_seen')))} | {hits} |"
            )
        lines.append("")
    lines.append("Outcome meanings: PROVEN requires handler hit, observed CDB LBA bytes, an LBA/current-state write or candidate PC, and packet/DMA state evidence. The fw source column means the data phase emitted an explicit 0x05EC firmware-LBA source decision event. PARTIAL means some firmware evidence was seen but the chain is incomplete. BLOCKED means this trace did not hit enough firmware-side evidence.")
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    ap = argparse.ArgumentParser(description="Validate GDR-8050L firmware LBA CPU/MMIO trace evidence")
    ap.add_argument("--trace-csv", required=True)
    ap.add_argument("--proof-pack", required=True)
    ap.add_argument("--out-json", default=None)
    ap.add_argument("--out-md", default=None)
    args = ap.parse_args()
    proof_pack = load_json_or_jsonc(args.proof_pack)
    report = validate_lba_trace(args.trace_csv, proof_pack)
    out_json = Path(args.out_json or str(Path(args.trace_csv).with_suffix(".lba_trace_proof.json")))
    out_md = Path(args.out_md or str(Path(args.trace_csv).with_suffix(".lba_trace_proof.md")))
    write_json(out_json, report)
    write_md(out_md, report)
    print(f"[OK] wrote LBA CPU trace proof JSON: {out_json}")
    print(f"[OK] wrote LBA CPU trace proof MD: {out_md}")
    bad = [f for f in report.get("findings") or [] if f.get("outcome") not in ("PROVEN", "PARTIAL")]
    return 2 if bad else 0


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