#!/usr/bin/env python3
"""Build a compact per-READ LBA proof report from GDR-8050L trace CSVs.

v190 proved one natural READ(10) can reach the firmware LBA assembler and
complete the persistent session.  v191 is the follow-up: run a small sweep of
READ LBAs and verify each READ command independently:

  * CDB LBA seen by the firmware-visible packet/FIFO path
  * firmware assembled-LBA value at 0x9001E541 / 0x05EC
  * DATA IN completion/submission for the same command window

Input can be a trace.csv, a route-scan artifact directory, or a .zip containing
one or more trace.csv files.  The script writes JSON and Markdown summaries.
"""
from __future__ import annotations

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

READ_OPS = {0x28: "READ10", 0xA8: "READ12", 0xBE: "READ_CD"}


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 norm_int(v: Any) -> int | None:
    try:
        return parse_int(v)
    except Exception:
        return None


def parse_bytes_hex(text: str | None) -> list[int]:
    if not text:
        return []
    out: list[int] = []
    for token in str(text).replace(",", " ").split():
        token = token.strip()
        if not token:
            continue
        try:
            out.append(int(token, 16) & 0xFF)
        except ValueError:
            try:
                out.append(parse_int(token) & 0xFF)
            except Exception:
                pass
    return out


def parse_notes_kv(notes: str | None) -> dict[str, str]:
    out: dict[str, str] = {}
    if not notes:
        return out
    for part in str(notes).split(";"):
        if "=" not in part:
            continue
        k, v = part.split("=", 1)
        out[k.strip()] = v.strip()
    return out


def cdb_lba_blocks(cdb: list[int]) -> tuple[int | None, int | None, str | None]:
    if not cdb:
        return None, None, None
    op = cdb[0]
    if op in (0x28, 0xBE) and len(cdb) >= 10:
        lba = ((cdb[2] << 24) | (cdb[3] << 16) | (cdb[4] << 8) | cdb[5]) & 0xFFFFFFFF
        blocks = ((cdb[7] << 8) | cdb[8]) & 0xFFFF
        if blocks == 0 and op == 0x28:
            blocks = 65536
        return lba, blocks, READ_OPS.get(op)
    if op == 0xA8 and len(cdb) >= 12:
        lba = ((cdb[2] << 24) | (cdb[3] << 16) | (cdb[4] << 8) | cdb[5]) & 0xFFFFFFFF
        blocks = ((cdb[6] << 24) | (cdb[7] << 16) | (cdb[8] << 8) | cdb[9]) & 0xFFFFFFFF
        return lba, blocks, READ_OPS.get(op)
    return None, None, None


@dataclass
class CommandProof:
    index: int
    tag: str
    opcode: int | None = None
    op_name: str | None = None
    cdb: list[int] = field(default_factory=list)
    cdb_lba: int | None = None
    blocks: int | None = None
    direction: str | None = None
    data_in_len_declared: int | None = None
    assembled_values: list[int] = field(default_factory=list)
    assembled_cycles: list[int] = field(default_factory=list)
    source_decisions: list[dict[str, Any]] = field(default_factory=list)
    data_in_source_bytes: int = 0
    data_in_first8: list[int] = field(default_factory=list)
    data_in_ready: bool = False
    data_in_submit: bool = False
    d108_complete: bool = False
    command_complete: bool = False
    d954_ready_reads: int = 0
    df38_fifo_reads: int = 0
    e541_hits: int = 0

    def is_read(self) -> bool:
        return self.opcode in READ_OPS

    def assembled_match(self) -> bool:
        return bool(self.assembled_values) and self.cdb_lba is not None and self.assembled_values[-1] == self.cdb_lba

    def expected_len(self) -> int | None:
        if self.data_in_len_declared is not None:
            return self.data_in_len_declared
        if self.blocks is not None:
            return int(self.blocks) * 2048
        return None

    def status(self) -> str:
        if not self.is_read():
            return "NON_READ"
        if not self.assembled_values:
            return "NO_ASSEMBLED_LBA"
        if not self.assembled_match():
            return "LBA_MISMATCH"
        exp = self.expected_len()
        if exp is not None and self.data_in_source_bytes not in (0, exp):
            return "PARTIAL_DATA_IN"
        if not self.d108_complete or not self.command_complete:
            return "NO_COMPLETION"
        return "PROVEN"

    def to_json(self) -> dict[str, Any]:
        return {
            "index": self.index,
            "tag": self.tag,
            "opcode": None if self.opcode is None else f"0x{self.opcode:02X}",
            "op_name": self.op_name,
            "cdb_hex": " ".join(f"{b:02X}" for b in self.cdb),
            "cdb_lba": self.cdb_lba,
            "cdb_lba_hex": None if self.cdb_lba is None else f"0x{self.cdb_lba:08X}",
            "blocks": self.blocks,
            "direction": self.direction,
            "data_in_len_declared": self.data_in_len_declared,
            "assembled_values": self.assembled_values,
            "assembled_values_hex": [f"0x{x:08X}" for x in self.assembled_values],
            "assembled_cycles": self.assembled_cycles,
            "assembled_match": self.assembled_match(),
            "source_decisions": self.source_decisions,
            "data_in_source_bytes": self.data_in_source_bytes,
            "data_in_first8_hex": " ".join(f"{b:02X}" for b in self.data_in_first8),
            "data_in_ready": self.data_in_ready,
            "data_in_submit": self.data_in_submit,
            "d108_complete": self.d108_complete,
            "command_complete": self.command_complete,
            "d954_ready_reads": self.d954_ready_reads,
            "df38_fifo_reads": self.df38_fifo_reads,
            "e541_hits": self.e541_hits,
            "status": self.status(),
        }


def load_trace_rows_from_file(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 trace_sources(paths: list[Path]) -> Iterable[tuple[str, list[dict[str, str]]]]:
    for path in paths:
        if path.is_dir():
            for trace in sorted(path.rglob("trace.csv")):
                yield str(trace), load_trace_rows_from_file(trace)
        elif path.suffix.lower() == ".zip":
            with zipfile.ZipFile(path) as z:
                for name in sorted(n for n in z.namelist() if n.endswith("trace.csv")):
                    text = z.read(name).decode("utf-8", errors="ignore")
                    yield f"{path}!{name}", [dict(r) for r in csv.DictReader(io.StringIO(text))]
        else:
            yield str(path), load_trace_rows_from_file(path)


def reduce_rows(source: str, rows: list[dict[str, str]]) -> dict[str, Any]:
    commands: list[CommandProof] = []
    current: CommandProof | None = None
    by_tag: dict[str, CommandProof] = {}

    def get_current_or_name(name: str | None) -> CommandProof | None:
        if name and name in by_tag:
            return by_tag[name]
        return current

    for r in rows:
        ev = r.get("event") or ""
        name = r.get("name") or ""
        pc = norm_int(r.get("pc"))
        cycle = parse_int(r.get("cycle"), 0)
        value = norm_int(r.get("value"))
        notes = parse_notes_kv(r.get("notes"))

        if ev == "gdr_session_command_arm":
            cdb = parse_bytes_hex(r.get("bytes_hex"))
            idx = parse_int(r.get("value"), len(commands))
            tag = name or f"cmd{idx:02d}"
            lba, blocks, opname = cdb_lba_blocks(cdb)
            current = CommandProof(index=idx, tag=tag, opcode=(cdb[0] if cdb else None), op_name=opname, cdb=cdb, cdb_lba=lba, blocks=blocks)
            current.direction = notes.get("direction")
            if "data_in_len" in notes:
                try:
                    current.data_in_len_declared = parse_int(notes.get("data_in_len"))
                except Exception:
                    pass
            commands.append(current)
            by_tag[tag] = current
            continue

        cmd = get_current_or_name(name)
        if cmd is None:
            continue

        if ev == "gdr_session_command_complete":
            target = by_tag.get(name, cmd)
            target.command_complete = True
        elif ev == "gdr_session_d108_hw_data_phase_complete":
            target = by_tag.get(name, cmd)
            target.d108_complete = True
        elif ev == "gdr_data_in_ready":
            target = by_tag.get(name, cmd)
            target.data_in_ready = True
        elif ev == "gdr_data_in_submit":
            target = by_tag.get(name, cmd)
            target.data_in_submit = True
        elif ev == "gdr_session_data_in_source":
            target = by_tag.get(name, cmd)
            target.data_in_source_bytes += 1
            if len(target.data_in_first8) < 8 and value is not None:
                target.data_in_first8.append(value & 0xFF)
        elif ev == "gdr_session_d954_ready_read":
            cmd.d954_ready_reads += 1
        elif ev == "gdr_session_df38_fifo_read":
            cmd.df38_fifo_reads += 1
        elif ev == "gdr_session_read_source_decision":
            target = by_tag.get(name, cmd)
            decision = {"pc": r.get("pc"), "cycle": cycle, "value": r.get("value"), "notes": r.get("notes")}
            target.source_decisions.append(decision)
        elif pc == 0x9001E541:
            cmd.e541_hits += 1
            # The break row at 0x9001E541 captures D0 immediately before the
            # mov d0,(0x05EC). That is the cleanest assembled-LBA value.
            d0 = norm_int(r.get("d0"))
            if d0 is not None and ev == "break":
                cmd.assembled_values.append(d0 & 0xFFFFFFFF)
                cmd.assembled_cycles.append(cycle)

    read_cmds = [c for c in commands if c.is_read()]
    return {
        "source": source,
        "rows": len(rows),
        "commands": [c.to_json() for c in commands],
        "read_commands": [c.to_json() for c in read_cmds],
        "read_count": len(read_cmds),
        "proven_read_count": sum(1 for c in read_cmds if c.status() == "PROVEN"),
        "failed_read_count": sum(1 for c in read_cmds if c.status() != "PROVEN"),
        "all_reads_proven": bool(read_cmds) and all(c.status() == "PROVEN" for c in read_cmds),
    }


def write_md(path: Path, report: dict[str, Any]) -> None:
    lines: list[str] = []
    lines.append("# GDR-8050L v191 READ/LBA sweep report")
    lines.append("")
    lines.append(f"inputs={len(report.get('traces') or [])}")
    lines.append(f"all_reads_proven={report.get('all_reads_proven')}")
    lines.append(f"proven_read_count={report.get('proven_read_count')} failed_read_count={report.get('failed_read_count')}")
    lines.append("")
    for trace in report.get("traces") or []:
        lines.append(f"## {trace.get('source')}")
        lines.append("")
        lines.append(f"rows={trace.get('rows')} read_count={trace.get('read_count')} proven={trace.get('proven_read_count')} failed={trace.get('failed_read_count')}")
        lines.append("")
        lines.append("| index | tag | op | CDB LBA | blocks | assembled 0x05EC | match | data bytes | complete | status |")
        lines.append("|---:|---|---|---:|---:|---:|---:|---:|---:|---|")
        for c in trace.get("read_commands") or []:
            asm = (c.get("assembled_values_hex") or [""])[-1] if c.get("assembled_values_hex") else ""
            lines.append(
                f"| {c.get('index')} | `{c.get('tag')}` | {c.get('op_name') or c.get('opcode')} | "
                f"{c.get('cdb_lba_hex') or ''} | {c.get('blocks') or ''} | {asm} | "
                f"{1 if c.get('assembled_match') else 0} | {c.get('data_in_source_bytes')} | "
                f"{1 if c.get('d108_complete') and c.get('command_complete') else 0} | {c.get('status')} |"
            )
        lines.append("")
        bad = [c for c in (trace.get("read_commands") or []) if c.get("status") != "PROVEN"]
        if bad:
            lines.append("### Non-proven READs")
            lines.append("")
            for c in bad:
                lines.append(f"- `{c.get('tag')}` status={c.get('status')} cdb={c.get('cdb_hex')} assembled={c.get('assembled_values_hex')} first8={c.get('data_in_first8_hex')}")
            lines.append("")
    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="Reduce GDR-8050L READ/LBA sweep traces")
    ap.add_argument("paths", nargs="+", type=Path, help="trace.csv, route-scan directory, or .zip")
    ap.add_argument("--json", dest="json_out", type=Path, default=Path("gdr_read_sweep_report.json"))
    ap.add_argument("--md", dest="md_out", type=Path, default=Path("gdr_read_sweep_report.md"))
    args = ap.parse_args()

    traces = [reduce_rows(src, rows) for src, rows in trace_sources(args.paths)]
    report = {
        "format": "hlds-gdr-read-sweep-report",
        "version": 1,
        "traces": traces,
        "proven_read_count": sum(t.get("proven_read_count", 0) for t in traces),
        "failed_read_count": sum(t.get("failed_read_count", 0) for t in traces),
        "all_reads_proven": bool(traces) and all(t.get("all_reads_proven") for t in traces),
    }
    args.json_out.parent.mkdir(parents=True, exist_ok=True)
    args.json_out.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    write_md(args.md_out, report)
    print(f"[OK] wrote {args.json_out}")
    print(f"[OK] wrote {args.md_out}")
    print(f"[READ-SWEEP] proven={report['proven_read_count']} failed={report['failed_read_count']} all={report['all_reads_proven']}")
    return 0


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