#!/usr/bin/env python3
"""Extract GDR-8050L packet/DF70 transaction artifacts from a trace CSV.

The emulator trace records packet frames and DF70 FIFO writes as watch events.
This helper turns those events into a compact JSON transcript plus optional raw
DF70 byte streams.  It is intentionally trace-driven, so it can be used on old
CSV files without rerunning the emulator.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import json
from collections import Counter
from pathlib import Path
from typing import Any

VERSION = "v116-known-cdb"

PACKET_PHASE_COPY_PCS = {
    "0x90024286", "0x9002428C", "0x90024292", "0x90024298",
    "0x9002429E", "0x900242A4", "0x900242AA", "0x900242B0",
}
PACKET_SERVICE_COPY_PCS = {
    "0x9002116C", "0x90021172", "0x90021178", "0x9002117E",
    "0x90021184", "0x9002118A", "0x90021190", "0x90021196",
}
DF70_START = 0x0000DF70
DF70_END = 0x0000DFAF  # current watch window is DF70:64

AUTH_LABELS: dict[int, str] = {
    0x000005B1: "primary packet/auth candidate flags",
    0x000005B4: "media/mech state flags",
    0x000005B6: "secondary media/auth flags",
    0x000005B7: "packet result flags",
    0x000005C6: "global media/error state flags",
    0x000005CD: "command completion/error flags",
    0x000005CE: "command completion/error flags",
    0x000005F0: "packet status byte",
    0x000006E1: "packet/media state",
    0x000006EC: "packet response bytes",
    0x000006ED: "packet response bytes",
    0x000006F1: "packet phase/result state",
    0x000006F2: "packet sub-state",
    0x000006F4: "packet callback/handler pointer",
    0x000006FC: "packet transfer timer/length",
    0x0000013B: "host status/completion flag",
    0x00000143: "drive state mirror",
    0x0000014A: "drive state mirror",
    0x0000037C: "auth/procedure state candidate",
    0x0000037D: "auth/procedure state candidate",
    0x0000037E: "auth/procedure timer/counter candidate",
    0x0000037F: "auth/procedure timer/counter candidate",
    0x0000DA00: "hardware/media status input",
    0x0000DA0C: "packet/media state",
    0x0000DA18: "packet mode flag",
    0x0000DA24: "packet/media state",
    0x0000DA28: "packet/media state",
    0x0000DA3C: "packet mode/status byte",
}


def parse_int(text: Any, default: int | None = None) -> int | None:
    try:
        return int(str(text).strip(), 0)
    except Exception:
        return default


def byte_from_row(row: dict[str, str]) -> int | None:
    v = parse_int(row.get("value"))
    if v is None:
        return None
    return v & 0xFF


def fmt_bytes(data: bytes | list[int], max_len: int | None = None) -> str:
    vals = list(data)
    if max_len is not None:
        vals = vals[:max_len]
    return " ".join(f"{b & 0xFF:02X}" for b in vals)


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


def copy_frames(rows: list[dict[str, str]], pcs: set[str], base_addr: int = 0x07E8) -> list[dict[str, Any]]:
    frames: list[dict[str, Any]] = []
    current: dict[int, int] = {}
    first_cycle: int | None = None
    last_cycle: int | None = None
    last_seq: int | None = None
    last_pc: str | None = None

    for row in rows:
        if row.get("event") != "watch_write" or row.get("pc") not in pcs:
            continue
        addr = parse_int(row.get("addr"))
        if addr is None or not (base_addr <= addr <= base_addr + 7):
            continue
        idx = addr - base_addr
        # A write to byte 0 after an already complete/partial frame starts a new frame.
        if idx == 0 and current:
            frames.append({
                "cycle_start": first_cycle,
                "cycle_end": last_cycle,
                "seq_end": last_seq,
                "pc_end": last_pc,
                "bytes": [current.get(i) for i in range(8)],
                "hex": " ".join("??" if current.get(i) is None else f"{current[i]:02X}" for i in range(8)),
            })
            current = {}
            first_cycle = None
        b = byte_from_row(row)
        if b is None:
            continue
        cyc = parse_int(row.get("cycle"), 0) or 0
        seq = parse_int(row.get("seq"), 0) or 0
        if first_cycle is None:
            first_cycle = cyc
        last_cycle = cyc
        last_seq = seq
        last_pc = row.get("pc")
        current[idx] = b
        if len(current) == 8:
            frames.append({
                "cycle_start": first_cycle,
                "cycle_end": last_cycle,
                "seq_end": last_seq,
                "pc_end": last_pc,
                "bytes": [current.get(i) for i in range(8)],
                "hex": " ".join(f"{current[i]:02X}" for i in range(8)),
            })
            current = {}
            first_cycle = None
            last_cycle = None
            last_seq = None
            last_pc = None
    if current:
        frames.append({
            "cycle_start": first_cycle,
            "cycle_end": last_cycle,
            "seq_end": last_seq,
            "pc_end": last_pc,
            "bytes": [current.get(i) for i in range(8)],
            "hex": " ".join("??" if current.get(i) is None else f"{current[i]:02X}" for i in range(8)),
        })
    return frames


def packet_start_cycle(rows: list[dict[str, str]]) -> int | None:
    candidates: list[int] = []
    for row in rows:
        pc = row.get("pc")
        if pc in {"0x9002425D", "0x90024286", "0x9002428C", "0x90024292"}:
            cyc = parse_int(row.get("cycle"))
            if cyc is not None:
                candidates.append(cyc)
    return min(candidates) if candidates else None


def df70_stream(rows: list[dict[str, str]], min_cycle: int = 0, max_cycle: int | None = None) -> tuple[bytes, list[dict[str, Any]]]:
    data = bytearray()
    events: list[dict[str, Any]] = []
    for row in rows:
        if row.get("event") != "watch_write":
            continue
        addr = parse_int(row.get("addr"))
        if addr is None or not (DF70_START <= addr <= DF70_END):
            continue
        cyc = parse_int(row.get("cycle"), 0) or 0
        if cyc < min_cycle:
            continue
        if max_cycle is not None and cyc > max_cycle:
            continue
        b = byte_from_row(row)
        if b is None:
            continue
        data.append(b)
        events.append({
            "seq": parse_int(row.get("seq"), 0),
            "cycle": cyc,
            "pc": row.get("pc"),
            "addr": row.get("addr"),
            "byte": b,
        })
    return bytes(data), events


def burst_summaries(events: list[dict[str, Any]], max_bursts: int = 32, gap_cycles: int = 50) -> tuple[int, list[dict[str, Any]]]:
    if not events:
        return 0, []
    bursts: list[list[dict[str, Any]]] = []
    cur: list[dict[str, Any]] = []
    last_cycle: int | None = None
    for ev in events:
        cyc = int(ev.get("cycle") or 0)
        if last_cycle is not None and cyc - last_cycle > gap_cycles:
            bursts.append(cur)
            cur = []
        cur.append(ev)
        last_cycle = cyc
    if cur:
        bursts.append(cur)
    out: list[dict[str, Any]] = []
    for i, burst in enumerate(bursts[:max_bursts]):
        bs = bytes(int(ev["byte"]) & 0xFF for ev in burst)
        out.append({
            "index": i,
            "cycle_start": burst[0].get("cycle"),
            "cycle_end": burst[-1].get("cycle"),
            "length": len(bs),
            "pc": burst[0].get("pc"),
            "first_addr": burst[0].get("addr"),
            "last_addr": burst[-1].get("addr"),
            "hex": fmt_bytes(bs, 64),
        })
    return len(bursts), out



def auth_candidate_info(rows: list[dict[str, str]], packet_start: int | None = None) -> dict[str, Any]:
    by_addr: dict[int, list[dict[str, str]]] = {}
    for row in rows:
        if not str(row.get("event", "")).startswith("watch"):
            continue
        addr = parse_int(row.get("addr"))
        if addr is None or addr not in AUTH_LABELS:
            continue
        by_addr.setdefault(addr, []).append(row)

    candidates: list[dict[str, Any]] = []
    for addr, evs in sorted(by_addr.items()):
        writes = [r for r in evs if r.get("event") == "watch_write"]
        reads = [r for r in evs if r.get("event") == "watch_read"]
        post_writes = [r for r in writes if packet_start is None or (parse_int(r.get("cycle"), 0) or 0) >= packet_start]
        write_events = []
        for row in writes[-32:]:
            write_events.append({
                "seq": parse_int(row.get("seq"), 0),
                "cycle": parse_int(row.get("cycle"), 0),
                "pc": row.get("pc"),
                "addr": row.get("addr"),
                "value": byte_from_row(row),
                "hex": "??" if byte_from_row(row) is None else f"{byte_from_row(row):02X}",
            })
        first = writes[0] if writes else (evs[0] if evs else {})
        last = writes[-1] if writes else (evs[-1] if evs else {})
        candidates.append({
            "addr": f"0x{addr:04X}",
            "label": AUTH_LABELS.get(addr, "auth/status candidate"),
            "reads": len(reads),
            "writes": len(writes),
            "post_packet_writes": len(post_writes),
            "first_cycle": parse_int(first.get("cycle"), 0),
            "first_pc": first.get("pc"),
            "last_cycle": parse_int(last.get("cycle"), 0),
            "last_pc": last.get("pc"),
            "last_value": byte_from_row(last) if last else None,
            "last_hex": "??" if not last or byte_from_row(last) is None else f"{byte_from_row(last):02X}",
            "write_events_tail": write_events,
        })
    ranked = sorted(candidates, key=lambda r: (r["post_packet_writes"], r["writes"], r["reads"]), reverse=True)
    return {
        "note": "0x05B1 bit 0x20 is the current strongest packet/auth candidate; v88 watches the surrounding flag/state bytes to prove or eliminate candidates.",
        "ranked": ranked,
    }


def read_capacity_candidate(data: bytes) -> dict[str, Any] | None:
    """Find a plausible READ CAPACITY(10)-style response in the DF70 stream.

    READ CAPACITY(10) is 8 bytes: last LBA and block length, both big-endian.
    The observed DF70 stream may contain short helper writes before the command
    payload, so v95 scans for a plausible 8-byte window rather than blindly
    trusting data[0:8].  The result remains candidate-only until the full ATAPI
    data phase/session model is proven.
    """
    if len(data) < 8:
        return None
    plausible_blocks = {512, 1024, 2048, 2056, 2064, 2324, 2336, 2352, 4096}

    def make(offset: int) -> dict[str, Any]:
        chunk = data[offset:offset + 8]
        lba = int.from_bytes(chunk[0:4], "big")
        block = int.from_bytes(chunk[4:8], "big")
        plausible = (block in plausible_blocks and 0 <= lba < 0x2000000)
        return {
            "candidate_only": True,
            "offset": offset,
            "first_8_hex": fmt_bytes(chunk),
            "last_lba_be": lba,
            "block_len_be": block,
            "sector_count_if_valid": lba + 1 if lba != 0xFFFFFFFF else None,
            "capacity_bytes_if_valid": (lba + 1) * block if block and lba != 0xFFFFFFFF else None,
            "plausible_window": plausible,
        }

    first = make(0)
    if first["plausible_window"]:
        return first
    # Search a small prefix.  Real READ CAPACITY output should be near the
    # beginning of the data-in burst; scanning farther risks false positives.
    limit = min(len(data) - 8, 256)
    for off in range(1, limit + 1):
        cand = make(off)
        if cand["plausible_window"]:
            cand["note"] = "plausible READ CAPACITY window found after helper/header bytes; data[0:8] was not used"
            cand["raw_stream_first_8_hex"] = fmt_bytes(data[:8])
            return cand
    first["note"] = "no plausible READ CAPACITY window found in first 256 bytes; data[0:8] shown for diagnostics"
    return first


def cdb_opcode_from_packet(packet_cdb: list[int]) -> int | None:
    return packet_cdb[0] if packet_cdb else None

def digest_info(data: bytes, events: list[dict[str, Any]]) -> dict[str, Any]:
    burst_count, bursts = burst_summaries(events)
    return {
        "length": len(data),
        "sha256": hashlib.sha256(data).hexdigest() if data else None,
        "first_64_hex": fmt_bytes(data, 64),
        "last_64_hex": fmt_bytes(data[-64:], None) if len(data) >= 64 else fmt_bytes(data),
        "burst_count": burst_count,
        "burst_samples": bursts,
    }




def _parse_known_cdb_text(text: str | None) -> list[int]:
    if not text:
        return []
    out: list[int] = []
    for part in str(text).replace(',', ' ').replace(';', ' ').split():
        part = part.strip()
        if not part:
            continue
        try:
            if part.lower().startswith(('0x', '-0x')):
                out.append(int(part, 0) & 0xFF)
            else:
                out.append(int(part, 16) & 0xFF)
        except Exception:
            continue
    return out


def _known_cdb_from_config(config_path: Path | None) -> list[int]:
    if not config_path:
        return []
    try:
        cfg = json.loads(config_path.read_text(encoding='utf-8'))
    except Exception:
        return []
    vals = cfg.get('packet_cdb') or cfg.get('host_cdb') or []
    if isinstance(vals, str):
        return _parse_known_cdb_text(vals)
    out: list[int] = []
    for item in vals:
        parsed = _parse_known_cdb_text(str(item))
        if parsed:
            out.extend(parsed[:1])
    return out

def build_transcript(csv_path: Path, df70_min_cycle: int | None = None, df70_max_cycle: int | None = None, known_cdb: list[int] | None = None) -> tuple[dict[str, Any], bytes, bytes]:
    rows = read_rows(csv_path)
    events = Counter(row.get("event", "") for row in rows)
    top_pc = Counter(row.get("pc", "") for row in rows).most_common(16)
    top_names = Counter(row.get("name", "") for row in rows).most_common(16)
    pkt_start = packet_start_cycle(rows)
    auto_min = pkt_start if df70_min_cycle is None else df70_min_cycle
    if auto_min is None:
        auto_min = 0

    all_df70, all_events = df70_stream(rows, 0, df70_max_cycle)
    post_df70, post_events = df70_stream(rows, auto_min, df70_max_cycle)
    phase_frames = copy_frames(rows, PACKET_PHASE_COPY_PCS)
    service_frames = copy_frames(rows, PACKET_SERVICE_COPY_PCS)

    # Reconstruct a likely 12-byte packet CDB from the first complete phase + service frames.
    packet_cdb: list[int] = []
    if phase_frames:
        b = phase_frames[0].get("bytes") or []
        packet_cdb.extend(int(x) for x in b[2:8] if x is not None)
    if service_frames:
        b = service_frames[0].get("bytes") or []
        packet_cdb.extend(int(x) for x in b[0:6] if x is not None)

    cap_candidate = None
    if packet_cdb and (packet_cdb[0] & 0xFF) == 0x25:
        cap_candidate = read_capacity_candidate(all_df70)
        if cap_candidate is not None:
            cap_candidate["source_stream"] = "df70_all"
    else:
        cap_candidate = read_capacity_candidate(all_df70)

    known = list(known_cdb or [])
    effective_cdb = known if known else packet_cdb
    cdb_source = "known/config" if known else "trace-reconstructed"

    transcript = {
        "tool": "extract_gdr_handshake.py",
        "version": VERSION,
        "source_csv": str(csv_path),
        "row_count": len(rows),
        "events": dict(events),
        "top_pc": top_pc,
        "top_names": top_names,
        "packet_start_cycle": pkt_start,
        "packet_phase_frames": phase_frames,
        "packet_service_frames": service_frames[:16],
        "reconstructed_packet_cdb_hex": fmt_bytes(packet_cdb),
        "reconstructed_packet_opcode": None if not packet_cdb else f"0x{packet_cdb[0] & 0xFF:02X}",
        "configured_packet_cdb_hex": fmt_bytes(known),
        "configured_packet_opcode": None if not known else f"0x{known[0] & 0xFF:02X}",
        "effective_packet_cdb_hex": fmt_bytes(effective_cdb),
        "effective_packet_opcode": None if not effective_cdb else f"0x{effective_cdb[0] & 0xFF:02X}",
        "cdb_source": cdb_source,
        "read_capacity10_candidate": cap_candidate,
        "df70_all": digest_info(all_df70, all_events),
        "df70_after_packet": digest_info(post_df70, post_events),
        "auth_candidates": auth_candidate_info(rows, pkt_start),
    }
    return transcript, all_df70, post_df70


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(description="Extract GDR packet/CDB/DF70 artifacts from a trace CSV")
    ap.add_argument("csv", type=Path, help="trace CSV from mn103s_emu.py")
    ap.add_argument("--out-json", type=Path, help="write transcript JSON; default <csv>.handshake.json")
    ap.add_argument("--out-df70", type=Path, help="write all DF70 bytes; default <csv>.df70_all.bin")
    ap.add_argument("--out-df70-after-packet", type=Path, help="write DF70 bytes after packet phase; default <csv>.df70_after_packet.bin")
    ap.add_argument("--df70-min-cycle", type=int, help="minimum cycle for after-packet stream; default auto-detected packet start")
    ap.add_argument("--df70-max-cycle", type=int, help="maximum cycle for DF70 streams")
    ap.add_argument("--known-cdb", help="authoritative full packet CDB hex/bytes; used when trace only exposes the service/tail frame")
    ap.add_argument("--config", type=Path, help="direct-config JSON to read packet_cdb/host_cdb from")
    ap.add_argument("--no-bin", action="store_true", help="only write JSON, not raw DF70 bin files")
    ns = ap.parse_args(argv[1:])

    csv_path = ns.csv.resolve()
    if not csv_path.exists():
        print(f"[-] missing trace CSV: {csv_path}")
        return 2
    known_cdb = _parse_known_cdb_text(ns.known_cdb) or _known_cdb_from_config(ns.config)
    transcript, all_df70, post_df70 = build_transcript(csv_path, ns.df70_min_cycle, ns.df70_max_cycle, known_cdb=known_cdb)
    out_json = ns.out_json or csv_path.with_suffix(csv_path.suffix + ".handshake.json")
    out_json.write_text(json.dumps(transcript, indent=2), encoding="utf-8")
    print(f"[OK] transcript: {out_json}")
    print(f"[CDB] {transcript.get('effective_packet_cdb_hex')} (source={transcript.get('cdb_source')})")
    recon = transcript.get('reconstructed_packet_cdb_hex')
    eff = transcript.get('effective_packet_cdb_hex')
    if recon and eff and recon != eff:
        print(f"[CDB-TRACE] reconstructed_tail={recon}")
    print(f"[DF70] all len={len(all_df70)} sha256={transcript['df70_all']['sha256']}")
    print(f"[DF70] after_packet len={len(post_df70)} sha256={transcript['df70_after_packet']['sha256']}")
    cap = transcript.get("read_capacity10_candidate") or {}
    if (transcript.get("effective_packet_opcode") or transcript.get("reconstructed_packet_opcode")) == "0x25" and cap:
        print(f"[READ-CAPACITY] candidate first8={cap.get('first_8_hex')} last_lba=0x{int(cap.get('last_lba_be') or 0):08X} block_len={cap.get('block_len_be')} capacity_bytes={cap.get('capacity_bytes_if_valid')}")
    for cand in transcript.get("auth_candidates", {}).get("ranked", [])[:8]:
        print(f"[AUTH] {cand['addr']} writes={cand['writes']} post={cand['post_packet_writes']} last={cand['last_hex']} pc={cand['last_pc']} label={cand['label']}")
    if not ns.no_bin:
        out_all = ns.out_df70 or csv_path.with_suffix(csv_path.suffix + ".df70_all.bin")
        out_post = ns.out_df70_after_packet or csv_path.with_suffix(csv_path.suffix + ".df70_after_packet.bin")
        out_all.write_bytes(all_df70)
        out_post.write_bytes(post_df70)
        print(f"[OK] df70_all: {out_all}")
        print(f"[OK] df70_after_packet: {out_post}")
    return 0


if __name__ == "__main__":
    import sys
    raise SystemExit(main(sys.argv))
