#!/usr/bin/env python3
"""Compare GDR-8050L auth/status candidates across extracted traces.

Runtime-trace-only: consumes extractor JSON files produced from CSV runs and
classifies low-RAM candidates.  It intentionally does not read Stage6/profile
metadata.

v94 keeps invariance analysis: a byte that changes the same way for every packet
CDB is demoted from "auth flag" to "generic packet/status flag" until a
challenge/data-out driven sequence proves otherwise.
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any

VERSION = "v100"


def load_json(path: Path) -> dict[str, Any]:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        raise SystemExit(f"[-] failed to read {path}: {exc}")


def cdb_hex(doc: dict[str, Any]) -> str:
    return str(doc.get("reconstructed_packet_cdb_hex") or doc.get("packet_cdb_hex") or "").strip()


def cdb_opcode(doc: dict[str, Any]) -> str:
    cdb = cdb_hex(doc)
    return cdb.split()[0].upper() if cdb.split() else "??"


def cdb_label(doc: dict[str, Any], fallback: str) -> str:
    cdb = cdb_hex(doc)
    if cdb:
        return f"{fallback} op=0x{cdb.split()[0].upper()} cdb={cdb}"
    return fallback


def candidate_rows(doc: dict[str, Any]) -> list[dict[str, Any]]:
    auth = doc.get("auth_candidates") or {}
    ranked = auth.get("ranked") or []
    return [r for r in ranked if isinstance(r, dict)]


def norm_addr(v: Any) -> str:
    return str(v or "").upper().replace("0X", "0x")


def classify(entries: list[dict[str, Any]], trace_count: int) -> tuple[str, str, int]:
    """Return classification, reason, auth_specific_score."""
    if not entries:
        return "no_events", "no entries", 0
    final_values = {str(e.get("last_hex")) for e in entries if e.get("last_hex") is not None}
    pcs = {str(e.get("last_pc")) for e in entries if e.get("last_pc")}
    post_by_trace = [int(e.get("post_packet_writes") or 0) for e in entries]
    post_positive = sum(1 for x in post_by_trace if x > 0)
    writes_positive = sum(1 for e in entries if int(e.get("writes") or 0) > 0)
    opcodes = {str(e.get("opcode") or "??") for e in entries}

    # Strong generic signal: seen in every trace, same final value, same writer PC,
    # and post-packet in every trace.  This is exactly what v90 showed for 0x05B1.
    if len(entries) == trace_count and trace_count > 1 and len(final_values) == 1 and len(pcs) == 1 and post_positive == trace_count:
        return (
            "generic_packet_status_invariant",
            "same post-packet write/final value/writer PC in every packet trace; not auth-specific in this run",
            0,
        )

    # Interesting: changes only for a subset of packet commands, or final values vary.
    if post_positive and (post_positive < trace_count or len(final_values) > 1 or len(pcs) > 1):
        score = post_positive * 100 + len(final_values) * 50 + len(pcs) * 10 + len(opcodes)
        return "auth_state_candidate_variant", "varies across CDBs or writer PCs", score

    if post_positive:
        return "packet_status_candidate", "post-packet writes observed but not enough variation to prove auth-specific", post_positive * 10

    if writes_positive:
        return "initialization_or_background_state", "only initialization/background writes observed", 1
    return "read_only_or_unmoved", "no writes observed", 0


def main(argv: list[str]) -> int:
    ap = argparse.ArgumentParser(description="Compare runtime auth/status candidate flags across GDR traces")
    ap.add_argument("json_files", nargs="+", type=Path, help="handshake/extractor JSON files")
    ap.add_argument("--out-json", type=Path, help="write machine-readable comparison JSON")
    ap.add_argument("--focus", default="0x05B1", help="candidate address to print first; default 0x05B1")
    ns = ap.parse_args(argv[1:])

    docs: list[tuple[Path, dict[str, Any]]] = [(p, load_json(p)) for p in ns.json_files]
    traces: list[dict[str, Any]] = []
    by_addr: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for path, doc in docs:
        label = cdb_label(doc, path.stem)
        op = cdb_opcode(doc)
        traces.append({
            "path": str(path),
            "name": path.stem,
            "label": label,
            "opcode": op,
            "cdb_hex": cdb_hex(doc),
            "row_count": doc.get("row_count"),
            "events": doc.get("events"),
        })
        for cand in candidate_rows(doc):
            addr = norm_addr(cand.get("addr"))
            if not addr:
                continue
            entry = {
                "trace": path.stem,
                "label": label,
                "opcode": op,
                "cdb_hex": cdb_hex(doc),
                "addr": addr,
                "candidate_label": cand.get("label"),
                "reads": int(cand.get("reads") or 0),
                "writes": int(cand.get("writes") or 0),
                "post_packet_writes": int(cand.get("post_packet_writes") or 0),
                "last_hex": cand.get("last_hex") or cand.get("last_value"),
                "last_value": cand.get("last_value"),
                "last_pc": cand.get("last_pc"),
                "last_cycle": cand.get("last_cycle"),
                "write_events_tail": cand.get("write_events_tail") or [],
            }
            by_addr[addr].append(entry)

    trace_count = len(traces)
    summary: list[dict[str, Any]] = []
    for addr, entries in sorted(by_addr.items()):
        final_values = sorted({str(e.get("last_hex")) for e in entries if e.get("last_hex") is not None})
        pcs = sorted({str(e.get("last_pc")) for e in entries if e.get("last_pc")})
        opcodes = sorted({str(e.get("opcode")) for e in entries})
        post_total = sum(int(e.get("post_packet_writes") or 0) for e in entries)
        write_total = sum(int(e.get("writes") or 0) for e in entries)
        read_total = sum(int(e.get("reads") or 0) for e in entries)
        kind, reason, auth_score = classify(entries, trace_count)
        generic_score = post_total * 100 + write_total * 10 + read_total + len(final_values) * 5
        summary.append({
            "addr": addr,
            "label": entries[0].get("candidate_label"),
            "trace_count": len(entries),
            "opcode_count": len(opcodes),
            "opcodes": opcodes,
            "write_total": write_total,
            "post_packet_write_total": post_total,
            "read_total": read_total,
            "final_values": final_values,
            "last_pcs": pcs,
            "classification": kind,
            "classification_reason": reason,
            "generic_packet_score": generic_score,
            "auth_specific_score": auth_score,
            "entries": entries,
        })

    ranked = sorted(summary, key=lambda r: (r["auth_specific_score"], r["post_packet_write_total"], r["write_total"]), reverse=True)
    generic_ranked = sorted(summary, key=lambda r: (r["generic_packet_score"], r["post_packet_write_total"], r["write_total"]), reverse=True)

    focus = norm_addr(ns.focus)
    print(f"[AUTH-COMPARE] version={VERSION} traces={len(traces)}")
    for t in traces:
        print(f"[TRACE] {Path(t['path']).name}: {t['label']}")

    focus_rows = [r for r in summary if r["addr"].upper() == focus.upper()]
    if focus_rows:
        r = focus_rows[0]
        tag = "DEMOTED" if r["classification"] == "generic_packet_status_invariant" else "FOCUS"
        print(
            f"[{tag}] {r['addr']} {r['label']} class={r['classification']} "
            f"post_total={r['post_packet_write_total']} final={','.join(r['final_values'])} "
            f"pcs={','.join(r['last_pcs'])} reason={r['classification_reason']}"
        )
        for e in r["entries"]:
            print(f"  - {e['trace']}: op=0x{e['opcode']} last={e['last_hex']} post={e['post_packet_writes']} writes={e['writes']} pc={e['last_pc']}")

    print("[AUTH-SPECIFIC-CANDIDATES]")
    any_variant = False
    for r in ranked[:12]:
        if r["auth_specific_score"] <= 10:
            continue
        any_variant = True
        print(
            f"  {r['addr']} {r['label']} class={r['classification']} auth_score={r['auth_specific_score']} "
            f"post={r['post_packet_write_total']} writes={r['write_total']} reads={r['read_total']} "
            f"final={','.join(r['final_values'])} pcs={','.join(r['last_pcs'])} opcodes={','.join(r['opcodes'])}"
        )
    if not any_variant:
        print("  none yet; current watched candidates are invariant or initialization-only")

    print("[GENERIC-PACKET-STATUS]")
    for r in generic_ranked[:8]:
        print(
            f"  {r['addr']} {r['label']} class={r['classification']} generic_score={r['generic_packet_score']} "
            f"post={r['post_packet_write_total']} final={','.join(r['final_values'])} pcs={','.join(r['last_pcs'])}"
        )

    out = {"tool": Path(__file__).name, "version": VERSION, "traces": traces, "ranked": ranked, "generic_ranked": generic_ranked}
    if ns.out_json:
        ns.out_json.write_text(json.dumps(out, indent=2), encoding="utf-8")
        print(f"[OK] wrote {ns.out_json}")
    return 0


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