#!/usr/bin/env python3
"""Combine corrected gate byte samples with branch-state pass/fail hypotheses."""
from __future__ import annotations

import argparse
import html
import json
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
SOURCE = "map1_01a"
TARGET = "map2_02d"
FAILED_GATE_PASS_MATRIX_GATE_IDS = [
    "current-selector-2-0-sample",
    "runtime-base-at-gate",
    "predecessor-state-persistence",
    "strict-source-hotspot",
]
GATE_PASS_MATRIX_MISSING_EVIDENCE = [
    "current selector 2:0 sample proving gate indices for this route",
    "runtime context+0xa8 base at 0x005428c4 and 0x005428cc",
    "predecessor 1:0 state persistence into the current 2:0 gate path",
    "strict map1_01a source hotspot",
]
GATE_PASS_MATRIX_EVIDENCE_REFS = [
    {
        "path": "out/save_selector_gate_sample_values.json",
        "description": "corrected public savedat gate-byte sample rows",
    },
    {
        "path": "out/save_selector_predecessor_state_effect.json",
        "description": "predecessor branch-state table hypothesis used for pass/fail evaluation",
    },
    {
        "path": "out/save_selector_gate_paths.json",
        "description": "current gate trace rows and branch/fallthrough addresses",
    },
]


def load_json(path: Path, fallback: Any) -> Any:
    if not path.exists():
        return fallback
    return json.loads(path.read_text(encoding="utf-8"))


def parse_hex(value: str | None) -> int | None:
    if not value:
        return None
    return int(value, 16)


def hex8(value: int | None) -> str | None:
    if value is None:
        return None
    return f"0x{value:02x}"


def gate_rows(gate_paths: list[dict]) -> list[dict]:
    if not gate_paths:
        return []
    row = gate_paths[0]
    traces = []
    for trace_name in ("firstGateTrace", "secondGateTrace"):
        trace = row.get(trace_name) or []
        if trace:
            traces.append(trace[0])
    gates = []
    for trace in traces:
        value = parse_hex(trace.get("valueHex"))
        stream_plus_1 = (value >> 8) & 0xFF if value is not None else None
        stream_plus_2 = (value >> 16) & 0xFF if value is not None else None
        gates.append({
            "gateVaHex": trace.get("vaHex"),
            "valueHex": trace.get("valueHex"),
            "selectionBufferOffset": stream_plus_2,
            "selectionBufferOffsetHex": hex8(stream_plus_2),
            "streamPlus1Hex": hex8(stream_plus_1),
            "stateTable": "primaryBranchState" if stream_plus_1 == 0 else "secondaryBranchState",
            "branchTargetHex": trace.get("branchTargetHex"),
            "fallthroughVaHex": trace.get("fallthroughVaHex"),
        })
    return gates


def table_hypotheses(predecessor_state_effect: dict) -> list[dict]:
    predecessor_table = predecessor_state_effect.get("secondaryBranchStateAfterFill")
    if not isinstance(predecessor_table, list) or len(predecessor_table) != 12:
        predecessor_table = [1, 1] + [0] * 10
    return [
        {
            "id": "predecessor_1_0_case0",
            "label": "predecessor 1:0 opcode 0x10 case 0",
            "stateTable": predecessor_table,
            "sourceEvidence": "save_selector_predecessor_state_effect",
        },
        {
            "id": "zero_table",
            "label": "all-zero secondaryBranchState",
            "stateTable": [0] * 12,
            "sourceEvidence": "negative control",
        },
    ]


def sample_groups(rows: list[dict]) -> dict[str, list[dict]]:
    grouped: dict[str, list[dict]] = {}
    for row in rows:
        sample_id = row.get("sampleId") or "-"
        grouped.setdefault(sample_id, []).append(row)
    for sample_id in grouped:
        grouped[sample_id].sort(key=lambda row: row.get("gateOffsetHex") or "")
    return dict(sorted(grouped.items()))


def evaluate_gate(row: dict, state_table: list[int]) -> dict:
    value = row.get("value")
    in_range = isinstance(value, int) and 0 <= value < len(state_table)
    state_value = state_table[value] if in_range else None
    return {
        "gateOffsetHex": row.get("gateOffsetHex"),
        "saveOffsetHex": row.get("saveOffsetHex"),
        "index": value,
        "indexHex": hex8(value) if isinstance(value, int) else None,
        "indexInRange": in_range,
        "stateValue": state_value,
        "fallsThrough": state_value == 1,
    }


def save_runtime_matrix(gate_sample_values: dict, hypotheses: list[dict]) -> list[dict]:
    grouped = sample_groups(gate_sample_values.get("saveRuntimeGateSampleRows") or [])
    rows = []
    for sample_id, sample_rows in grouped.items():
        sample = sample_rows[0]
        for hypothesis in hypotheses:
            gate_results = [
                evaluate_gate(row, hypothesis["stateTable"])
                for row in sample_rows
            ]
            rows.append({
                "sampleId": sample_id,
                "selector": sample.get("selector"),
                "tile": sample.get("tile"),
                "coversCurrentFrontierSelector": sample.get("coversCurrentFrontierSelector"),
                "hypothesisId": hypothesis["id"],
                "hypothesisLabel": hypothesis["label"],
                "gateResults": gate_results,
                "allGateIndicesInRange": all(row["indexInRange"] for row in gate_results),
                "allGatesFallThrough": all(row["fallsThrough"] for row in gate_results),
            })
    return rows


def party_slot_summary(gate_sample_values: dict) -> dict:
    rows = gate_sample_values.get("partySlotStatSampleRows") or []
    in_range = [row for row in rows if row.get("inExpectedIndexRange")]
    distinct = sorted({row.get("valueHex") for row in rows if row.get("valueHex")})
    return {
        "sampleValueCount": len(rows),
        "inRangeSampleValueCount": len(in_range),
        "outOfRangeSampleValueCount": len(rows) - len(in_range),
        "distinctValuesHex": distinct,
        "usableAsGateIndexEvidence": len(in_range) > 0,
    }


def build_summary(
    gate_sample_values: dict | None = None,
    predecessor_state_effect: dict | None = None,
    gate_paths: list[dict] | None = None,
) -> dict:
    gate_sample_values = gate_sample_values if gate_sample_values is not None else load_json(
        OUT / "save_selector_gate_sample_values.json",
        {},
    )
    predecessor_state_effect = predecessor_state_effect if predecessor_state_effect is not None else load_json(
        OUT / "save_selector_predecessor_state_effect.json",
        {},
    )
    gate_paths = gate_paths if gate_paths is not None else load_json(OUT / "save_selector_gate_paths.json", [])
    hypotheses = table_hypotheses(predecessor_state_effect)
    matrix = save_runtime_matrix(gate_sample_values, hypotheses)
    predecessor_rows = [row for row in matrix if row["hypothesisId"] == "predecessor_1_0_case0"]
    zero_rows = [row for row in matrix if row["hypothesisId"] == "zero_table"]
    party_summary = party_slot_summary(gate_sample_values)
    conclusion = (
        "With the corrected save/runtime mapping, public samples put both gate indices at 0. Under the strongest "
        "predecessor 1:0 hypothesis, secondaryBranchState[0] is 1, so both current control gates "
        "0x005428c4 and 0x005428cc would fall through for every public sample. Under an all-zero table they would "
        "not. This narrows the gate problem but does not promote map1_01a -> map2_02d: the public samples do not "
        "cover selector 2:0, the runtime context+0xa8 base at the gates is still unproven, and there is still no "
        "strict map1_01a source hotspot."
    )
    return {
        "source": SOURCE,
        "target": TARGET,
        "gateRows": gate_rows(gate_paths),
        "currentFrontierSelector": gate_sample_values.get("currentFrontierSelector"),
        "currentFrontierSampleCovered": gate_sample_values.get("currentFrontierSampleCovered"),
        "saveRuntimeGateDistinctValuesHex": gate_sample_values.get("saveRuntimeGateDistinctValuesHex") or [],
        "tableHypotheses": hypotheses,
        "saveRuntimePassMatrix": matrix,
        "saveRuntimePredecessorAllGatePassSampleCount": sum(1 for row in predecessor_rows if row["allGatesFallThrough"]),
        "saveRuntimePredecessorSampleCount": len(predecessor_rows),
        "saveRuntimeZeroTableAllGatePassSampleCount": sum(1 for row in zero_rows if row["allGatesFallThrough"]),
        "saveRuntimeZeroTableSampleCount": len(zero_rows),
        "partySlotStatSummary": party_summary,
        "runtimeBaseProofRequired": True,
        "predecessorPersistenceProofRequired": True,
        "strictHotspotProofRequired": True,
        "proofFound": False,
        "gatePassMatrixProofFound": False,
        "failedGatePassMatrixGateIds": FAILED_GATE_PASS_MATRIX_GATE_IDS,
        "missingEvidence": GATE_PASS_MATRIX_MISSING_EVIDENCE,
        "evidenceRefs": GATE_PASS_MATRIX_EVIDENCE_REFS,
        "evidenceRefCount": len(GATE_PASS_MATRIX_EVIDENCE_REFS),
        "promotionStatus": "blocked",
        "conclusion": conclusion,
    }


def gate_result_text(results: list[dict]) -> str:
    pieces = []
    for row in results:
        state = "-" if row["stateValue"] is None else str(row["stateValue"])
        pieces.append(
            f"{row['gateOffsetHex']} idx {row['indexHex'] or '-'} -> state {state} -> "
            f"{'fallthrough' if row['fallsThrough'] else 'branch'}"
        )
    return "; ".join(pieces)


def markdown(summary: dict) -> str:
    party = summary["partySlotStatSummary"]
    lines = [
        "# Save Selector Gate Pass Matrix",
        "",
        f"- route: {summary['source']} -> {summary['target']}",
        f"- current frontier selector: `{summary['currentFrontierSelector']}`",
        f"- current frontier sample covered: {summary['currentFrontierSampleCovered']}",
        f"- save/runtime distinct gate values: {', '.join(f'`{item}`' for item in summary['saveRuntimeGateDistinctValuesHex'])}",
        f"- predecessor pass samples: {summary['saveRuntimePredecessorAllGatePassSampleCount']} / {summary['saveRuntimePredecessorSampleCount']}",
        f"- zero-table pass samples: {summary['saveRuntimeZeroTableAllGatePassSampleCount']} / {summary['saveRuntimeZeroTableSampleCount']}",
        f"- party-slot usable in-range values: {party['inRangeSampleValueCount']} / {party['sampleValueCount']}",
        f"- runtime base proof required: {summary['runtimeBaseProofRequired']}",
        f"- predecessor persistence proof required: {summary['predecessorPersistenceProofRequired']}",
        f"- strict hotspot proof required: {summary['strictHotspotProofRequired']}",
        f"- proofFound: `{summary['proofFound']}`",
        f"- gatePassMatrixProofFound: `{summary['gatePassMatrixProofFound']}`",
        f"- promotion status: `{summary['promotionStatus']}`",
        "",
        summary["conclusion"],
        "",
        "## Failed Gates",
        "",
    ]
    lines.extend(f"- `{item}`" for item in summary["failedGatePassMatrixGateIds"])
    lines.extend([
        "",
        "## Missing Evidence",
        "",
    ])
    lines.extend(f"- {item}" for item in summary["missingEvidence"])
    lines.extend([
        "",
        "## Evidence Refs",
        "",
    ])
    lines.extend(
        f"- `{row['path']}`: {row['description']}"
        for row in summary["evidenceRefs"]
    )
    lines.extend([
        "",
        "## Current Gates",
        "",
        "| gate | value | selection offset | table | branch target | fallthrough |",
        "| --- | --- | --- | --- | --- | --- |",
    ])
    for row in summary["gateRows"]:
        lines.append(
            f"| `{row['gateVaHex']}` | `{row['valueHex']}` | `{row['selectionBufferOffsetHex']}` | "
            f"{row['stateTable']} | `{row['branchTargetHex']}` | `{row['fallthroughVaHex']}` |"
        )
    lines.extend([
        "",
        "## Save/Runtime Matrix",
        "",
        "| sample | selector | hypothesis | gate results | all gates fall through |",
        "| --- | --- | --- | --- | --- |",
    ])
    for row in summary["saveRuntimePassMatrix"]:
        lines.append(
            f"| {row['sampleId']} | `{row['selector']}` | {row['hypothesisLabel']} | "
            f"{gate_result_text(row['gateResults'])} | {row['allGatesFallThrough']} |"
        )
    lines.extend([
        "",
        "## Party Slot Stat Summary",
        "",
        f"- sample values: {party['sampleValueCount']}",
        f"- in range: {party['inRangeSampleValueCount']}",
        f"- out of range: {party['outOfRangeSampleValueCount']}",
        f"- distinct values: {', '.join(f'`{item}`' for item in party['distinctValuesHex'])}",
        "",
    ])
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    failed_gates = "".join(
        f"<li><code>{html.escape(item)}</code></li>"
        for item in summary["failedGatePassMatrixGateIds"]
    )
    missing_evidence = "".join(
        f"<li>{html.escape(item)}</li>"
        for item in summary["missingEvidence"]
    )
    evidence_refs = "".join(
        f"<li><code>{html.escape(row['path'])}</code>: {html.escape(row['description'])}</li>"
        for row in summary["evidenceRefs"]
    )
    gate_rows_html = []
    for row in summary["gateRows"]:
        gate_rows_html.append(
            "<tr>"
            f"<td><code>{html.escape(row['gateVaHex'] or '-')}</code></td>"
            f"<td><code>{html.escape(row['valueHex'] or '-')}</code></td>"
            f"<td><code>{html.escape(row['selectionBufferOffsetHex'] or '-')}</code></td>"
            f"<td>{html.escape(row['stateTable'])}</td>"
            f"<td><code>{html.escape(row['branchTargetHex'] or '-')}</code></td>"
            f"<td><code>{html.escape(row['fallthroughVaHex'] or '-')}</code></td>"
            "</tr>"
        )
    matrix_rows_html = []
    for row in summary["saveRuntimePassMatrix"]:
        matrix_rows_html.append(
            "<tr>"
            f"<td>{html.escape(row['sampleId'])}</td>"
            f"<td><code>{html.escape(row['selector'])}</code></td>"
            f"<td>{html.escape(row['hypothesisLabel'])}</td>"
            f"<td>{html.escape(gate_result_text(row['gateResults']))}</td>"
            f"<td>{row['allGatesFallThrough']}</td>"
            "</tr>"
        )
    party = summary["partySlotStatSummary"]
    return "\n".join([
        "<!doctype html>",
        '<html lang="en">',
        "<head>",
        '  <meta charset="utf-8">',
        '  <meta name="viewport" content="width=device-width, initial-scale=1">',
        "  <title>Save Selector Gate Pass Matrix</title>",
        "  <style>",
        "    :root { color-scheme: dark; font-family: system-ui, sans-serif; background: #101010; color: #eee; }",
        "    body { margin: 0; padding: 24px; }",
        "    h1 { margin: 0 0 8px; font-size: 24px; }",
        "    h2 { margin: 24px 0 8px; font-size: 18px; }",
        "    p { max-width: 1120px; color: #bbb; line-height: 1.45; }",
        "    table { width: 100%; border-collapse: collapse; margin: 12px 0 20px; font-size: 13px; }",
        "    th, td { border-bottom: 1px solid #303030; padding: 7px 8px; text-align: left; vertical-align: top; }",
        "    th { position: sticky; top: 0; background: #181818; color: #ddd; }",
        "    code { color: #d7f0ff; }",
        "  </style>",
        "</head>",
        "<body>",
        "  <h1>Save Selector Gate Pass Matrix</h1>",
        f"  <p>route {html.escape(summary['source'])} -&gt; {html.escape(summary['target'])}; "
        f"predecessor pass samples {summary['saveRuntimePredecessorAllGatePassSampleCount']} / {summary['saveRuntimePredecessorSampleCount']}; "
        f"zero-table pass samples {summary['saveRuntimeZeroTableAllGatePassSampleCount']} / {summary['saveRuntimeZeroTableSampleCount']}; "
        f"party-slot in-range values {party['inRangeSampleValueCount']} / {party['sampleValueCount']}; "
        f"runtime base proof required: {summary['runtimeBaseProofRequired']}; "
        f"predecessor persistence proof required: {summary['predecessorPersistenceProofRequired']}; "
        f"strict hotspot proof required: {summary['strictHotspotProofRequired']}; "
        f"proofFound <code>{summary['proofFound']}</code>; "
        f"promotion status <code>{html.escape(summary['promotionStatus'])}</code>.</p>",
        f"  <p>{html.escape(summary['conclusion'])}</p>",
        "  <h2>Failed Gates</h2>",
        f"  <ul>{failed_gates}</ul>",
        "  <h2>Missing Evidence</h2>",
        f"  <ul>{missing_evidence}</ul>",
        "  <h2>Evidence Refs</h2>",
        f"  <ul>{evidence_refs}</ul>",
        "  <h2>Current Gates</h2>",
        "  <table><thead><tr><th>gate</th><th>value</th><th>selection offset</th><th>table</th><th>branch target</th><th>fallthrough</th></tr></thead><tbody>",
        *gate_rows_html,
        "  </tbody></table>",
        "  <h2>Save/Runtime Matrix</h2>",
        "  <table><thead><tr><th>sample</th><th>selector</th><th>hypothesis</th><th>gate results</th><th>all gates fall through</th></tr></thead><tbody>",
        *matrix_rows_html,
        "  </tbody></table>",
        "  <h2>Party Slot Stat Summary</h2>",
        f"  <p>sample values {party['sampleValueCount']}; in range {party['inRangeSampleValueCount']}; out of range {party['outOfRangeSampleValueCount']}; distinct values {html.escape(', '.join(party['distinctValuesHex']))}.</p>",
        "</body>",
        "</html>",
        "",
    ])


def write_outputs(summary: dict, out_dir: Path = OUT) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_gate_pass_matrix.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (out_dir / "save_selector_gate_pass_matrix.html").write_text(html_page(summary), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gate-sample-values", type=Path, default=OUT / "save_selector_gate_sample_values.json")
    parser.add_argument("--predecessor-state-effect", type=Path, default=OUT / "save_selector_predecessor_state_effect.json")
    parser.add_argument("--gate-paths", type=Path, default=OUT / "save_selector_gate_paths.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    summary = build_summary(
        load_json(args.gate_sample_values, {}),
        load_json(args.predecessor_state_effect, {}),
        load_json(args.gate_paths, []),
    )
    write_outputs(summary, args.out_dir)
    print(f"wrote save selector gate pass matrix -> {args.out_dir / 'save_selector_gate_pass_matrix.html'}")


if __name__ == "__main__":
    main()