#!/usr/bin/env python3
"""Summarize the selector leaf table around the current 2:0 frontier."""
from __future__ import annotations

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

sys.path.insert(0, str(Path(__file__).resolve().parent))

from probe_exe_scene_tables import find_cns_strings, read_sections, va_to_offset
from summarize_script_handler_table import handler_for_opcode, section_name_for_va
from summarize_save_selector_stream_traces import trace_stream


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"

CURRENT_ROOT = 0x00540714
TABLE_START = 0x005429DC
WINDOW_START = 0x005429A8
WINDOW_END = 0x00542A10
WRAPPER_LEAF = 0x00542A04
FRONTIER_LEAF = 0x00542AE8
FRONTIER_READER = 0x00542B0C

EVIDENCE_REFS = [
    {
        "path": "out/save_scene_selector_references.json",
        "fields": [
            "label",
            "kind",
            "resource",
            "pathHex",
        ],
    },
    {
        "path": "out/save_selector_leaf_streams.json",
        "fields": [
            "leafPointerHex",
            "fieldRecords",
        ],
    },
    {
        "path": "out/save_selector_stream_traces.json",
        "fields": [
            "streamVaHex",
            "trace",
        ],
    },
    {
        "path": "out/script_handler_table.json",
        "fields": [
            "handlerTableVaHex",
            "entries",
            "codeHandlerCount",
            "defaultHandlerCount",
        ],
    },
]

LEAF_TABLE_MISSING_EVIDENCE_BY_GATE = {
    "runtime-selector-index": (
        "runtime selector/index that chooses 0x00542a04 or 0x00542ae8 from the table window"
    ),
    "wrapper-frontier-execution": (
        "wrapper 0x00542a04 execution into 0x00542ae8 on the normal route"
    ),
    "strict-hotspot": "strict map1_01a source coordinate or hotspot",
}


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


def hex32(value: int) -> str:
    return f"0x{value:08x}"


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


def dword_at(exe: bytes, sections: list[dict], va: int) -> int | None:
    offset = va_to_offset(sections, va)
    if offset is None or offset + 4 > len(exe):
        return None
    return struct.unpack_from("<I", exe, offset)[0]


def refs_to_value(exe: bytes, sections: list[dict], value: int) -> list[dict]:
    needle = struct.pack("<I", value)
    refs = []
    for section in sections:
        start = section["raw"]
        end = start + section["raw_size"]
        data = exe[start:end]
        pos = data.find(needle)
        while pos >= 0:
            refs.append({
                "section": section["name"],
                "refVaHex": hex32(section["va"] + pos),
            })
            pos = data.find(needle, pos + 1)
    return refs


def field_maps_for_leaf(leaf_streams: list[dict], leaf_hex: str) -> list[str]:
    row = next((item for item in leaf_streams if item.get("leafPointerHex") == leaf_hex), None)
    if not row:
        return []
    return [record.get("map") for record in row.get("fieldRecords") or [] if record.get("map")]


def trace_contains_reader(stream_traces: list[dict], leaf_hex: str) -> bool:
    row = next((item for item in stream_traces if item.get("streamVaHex") == leaf_hex), None)
    if not row:
        return False
    return any(item.get("vaHex") == hex32(FRONTIER_READER) for item in row.get("trace") or [])


def row_at(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    va: int,
    leaf_streams: list[dict],
    stream_traces: list[dict],
) -> dict:
    value = dword_at(exe, sections, va)
    row = {
        "vaHex": hex32(va),
        "valueHex": hex32(value) if value is not None else None,
        "insideRootTableWindow": TABLE_START <= va < WRAPPER_LEAF,
    }
    if value is None:
        row["kind"] = "unreadable"
        return row
    opcode = value & 0xFF
    handler = handler_for_opcode(exe, sections, opcode)
    section = section_name_for_va(sections, value)
    row.update({
        "lowByteHex": f"0x{opcode:02x}",
        "lowByteHandlerVaHex": handler.get("handlerVaHex"),
        "lowByteHandlerSection": handler.get("handlerSection"),
        "lowByteIsCodeHandler": handler.get("isCodeHandler"),
    })
    value_hex = hex32(value)
    leaf_maps = field_maps_for_leaf(leaf_streams, value_hex)
    if value == WRAPPER_LEAF:
        row["kind"] = "wrapperLeafPointer"
    elif value == FRONTIER_LEAF:
        row["kind"] = "frontierLeafPointer"
    elif leaf_maps:
        row["kind"] = "leafPointer"
    elif value in strings:
        row["kind"] = "cns"
        row["cns"] = strings[value]
    elif section:
        row["kind"] = "pointer"
        row["targetVaHex"] = value_hex
        row["targetSection"] = section
    else:
        row["kind"] = "scalar"
    if leaf_maps:
        row["leafPointerHex"] = value_hex
        row["leafFieldMaps"] = leaf_maps
        row["traceContainsFrontierReader"] = trace_contains_reader(stream_traces, value_hex)
    return row


def compact_trace(exe: bytes, sections: list[dict], strings: dict[int, str], start: int) -> list[dict]:
    rows = []
    for item in trace_stream(exe, sections, strings, start, 8):
        rows.append({
            "step": item.get("step"),
            "vaHex": item.get("vaHex"),
            "valueHex": item.get("valueHex"),
            "opcodeHex": item.get("opcodeHex"),
            "handlerVaHex": item.get("handlerVaHex"),
            "stopReason": item.get("stopReason"),
            "branchTargetHex": item.get("branchTargetHex"),
            "fallthroughVaHex": item.get("fallthroughVaHex"),
        })
        if item.get("stopReason"):
            break
    return rows


def build_summary(
    exe: bytes,
    selector_references: list[dict],
    leaf_streams: list[dict],
    stream_traces: list[dict],
) -> dict:
    sections = read_sections(exe)
    strings = find_cns_strings(exe, sections)
    table_rows = [
        row_at(exe, sections, strings, va, leaf_streams, stream_traces)
        for va in range(WINDOW_START, WINDOW_END + 1, 4)
    ]
    leaf_refs = [
        row for row in table_rows
        if row.get("kind") in {"leafPointer", "frontierLeafPointer", "wrapperLeafPointer"}
    ]
    frontier_ref = next((row for row in leaf_refs if row.get("kind") == "frontierLeafPointer"), None)
    wrapper_trace = compact_trace(exe, sections, strings, WRAPPER_LEAF)
    frontier_trace = compact_trace(exe, sections, strings, FRONTIER_LEAF)
    selector_paths = [
        row for row in selector_references
        if row.get("label") == "2:0"
        and row.get("kind") == "fieldMap"
        and row.get("resource") in {"map1_01a", "map2_02d"}
        and "0x005429dc" in (row.get("pathHex") or [])
    ]
    direct_leaf_refs = {
        hex32(value): refs_to_value(exe, sections, value)
        for value in [0x00542AC0, 0x00542AD4, FRONTIER_LEAF, WRAPPER_LEAF]
    }
    runtime_selection_proven = False
    strict_hotspot_found = False
    failed_leaf_table_gate_ids = []
    if not runtime_selection_proven:
        failed_leaf_table_gate_ids.append("runtime-selector-index")
    if not trace_contains_reader(stream_traces, hex32(WRAPPER_LEAF)):
        failed_leaf_table_gate_ids.append("wrapper-frontier-execution")
    if not strict_hotspot_found:
        failed_leaf_table_gate_ids.append("strict-hotspot")
    missing_evidence = [
        LEAF_TABLE_MISSING_EVIDENCE_BY_GATE.get(gate_id, gate_id)
        for gate_id in failed_leaf_table_gate_ids
    ]
    conclusion = (
        "The current selector root points at the 0x005429dc leaf table window. That window directly references "
        "0x00542ac0 and 0x00542ad4, and it reaches the reader-bearing frontier leaf 0x00542ae8 through a small "
        "wrapper at 0x00542a04 whose second dword points at 0x00542ae8. This strengthens the evidence that the "
        "frontier leaf belongs to the current selector table, but the wrapper trace still stops on a data-backed "
        "low byte and does not prove runtime selection of 0x00542ae8. Strict source hotspot evidence is still missing."
    )
    return {
        "selector": "2:0",
        "rootHex": hex32(CURRENT_ROOT),
        "rootTablePointerHex": hex32(TABLE_START),
        "tableWindowHex": f"{hex32(WINDOW_START)}..{hex32(WINDOW_END)}",
        "evidenceRefs": EVIDENCE_REFS,
        "evidenceRefCount": len(EVIDENCE_REFS),
        "wrapperLeafHex": hex32(WRAPPER_LEAF),
        "frontierLeafHex": hex32(FRONTIER_LEAF),
        "frontierReaderHex": hex32(FRONTIER_READER),
        "tableRows": table_rows,
        "leafRefCount": len(leaf_refs),
        "leafRefs": leaf_refs,
        "frontierLeafRefVaHex": frontier_ref.get("vaHex") if frontier_ref else None,
        "frontierLeafRefIsDirectRootTableEntry": bool(frontier_ref and frontier_ref.get("insideRootTableWindow")),
        "wrapperTrace": wrapper_trace,
        "frontierTrace": frontier_trace,
        "selectorPathCount": len(selector_paths),
        "selectorPaths": selector_paths[:12],
        "directLeafRefs": direct_leaf_refs,
        "runtimeSelectionProven": runtime_selection_proven,
        "proofFound": runtime_selection_proven,
        "failedLeafTableGateIds": failed_leaf_table_gate_ids,
        "missingEvidence": missing_evidence,
        "strictHotspotFound": strict_hotspot_found,
        "promotionStatus": "blocked",
        "remainingProofs": [
            "decode the runtime selector/index that chooses 0x00542a04 or 0x00542ae8 from the table window",
            "prove the wrapper at 0x00542a04 executes into 0x00542ae8 in the normal route",
            "find strict map1_01a source coordinate or hotspot",
        ],
        "conclusion": conclusion,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Leaf Table Context",
        "",
        f"- selector: `{summary['selector']}` root `{summary['rootHex']}`",
        f"- root table pointer: `{summary['rootTablePointerHex']}`",
        f"- table window: `{summary['tableWindowHex']}`",
        f"- wrapper leaf: `{summary['wrapperLeafHex']}`",
        f"- frontier leaf: `{summary['frontierLeafHex']}`",
        f"- frontier reader: `{summary['frontierReaderHex']}`",
        f"- leaf refs in window: {summary['leafRefCount']}",
        f"- frontier leaf ref: `{summary['frontierLeafRefVaHex']}`",
        f"- frontier leaf is direct root table entry: {summary['frontierLeafRefIsDirectRootTableEntry']}",
        f"- selector path rows for map1_01a/map2_02d: {summary['selectorPathCount']}",
        f"- runtime selection proven: {summary['runtimeSelectionProven']}",
        f"- proof found: {summary['proofFound']}",
        f"- failed leaf-table gates: `{', '.join(summary.get('failedLeafTableGateIds') or []) or '-'}`",
        f"- missing evidence count: {len(summary.get('missingEvidence') or [])}",
        f"- strict hotspot found: {summary['strictHotspotFound']}",
        f"- promotion status: {summary['promotionStatus']}",
        f"- evidence refs: `{summary['evidenceRefCount']}`",
        "",
        summary["conclusion"],
        "",
    ]
    lines.extend(["## Missing Evidence", ""])
    lines.extend(f"- {item}" for item in summary.get("missingEvidence") or [])
    lines.extend(["", "## Evidence Refs", ""])
    for ref in summary.get("evidenceRefs") or []:
        lines.append(
            f"- `{ref.get('path')}`: {', '.join(ref.get('fields') or [])}"
        )
    lines.extend([
        "",
        "## Table Window",
        "",
        "| va | value | kind | leaf maps | reader? | handler |",
        "| --- | --- | --- | --- | --- | --- |",
    ])
    for row in summary["tableRows"]:
        lines.append(
            f"| `{row['vaHex']}` | `{row['valueHex']}` | {row.get('kind')} | "
            f"{', '.join(row.get('leafFieldMaps') or []) or '-'} | "
            f"{row.get('traceContainsFrontierReader', '-')} | "
            f"`{row.get('lowByteHandlerVaHex') or '-'}` {row.get('lowByteHandlerSection') or '-'} |"
        )
    lines.extend(["", "## Wrapper Trace", "", "| step | va | value | opcode | handler | stop |", "| ---: | --- | --- | --- | --- | --- |"])
    for row in summary["wrapperTrace"]:
        lines.append(
            f"| {row.get('step')} | `{row.get('vaHex')}` | `{row.get('valueHex')}` | "
            f"`{row.get('opcodeHex')}` | `{row.get('handlerVaHex')}` | {row.get('stopReason') or '-'} |"
        )
    lines.extend(["", "## Frontier Trace", "", "| step | va | value | opcode | handler | branch target | fallthrough | stop |", "| ---: | --- | --- | --- | --- | --- | --- | --- |"])
    for row in summary["frontierTrace"]:
        lines.append(
            f"| {row.get('step')} | `{row.get('vaHex')}` | `{row.get('valueHex')}` | "
            f"`{row.get('opcodeHex')}` | `{row.get('handlerVaHex')}` | "
            f"`{row.get('branchTargetHex') or '-'}` | `{row.get('fallthroughVaHex') or '-'}` | "
            f"{row.get('stopReason') or '-'} |"
        )
    lines.extend(["", "## Remaining Proofs", ""])
    lines.extend(f"- {item}" for item in summary["remainingProofs"])
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    evidence_refs = "".join(
        "<li>"
        f"<code>{html.escape(str(ref.get('path')))}</code>: "
        f"{html.escape(', '.join(ref.get('fields') or []))}"
        "</li>"
        for ref in summary.get("evidenceRefs") or []
    )
    table_rows = "\n".join(
        "<tr>"
        f"<td><code>{html.escape(row['vaHex'])}</code></td>"
        f"<td><code>{html.escape(str(row['valueHex']))}</code></td>"
        f"<td>{html.escape(str(row.get('kind')))}</td>"
        f"<td>{html.escape(', '.join(row.get('leafFieldMaps') or []) or '-')}</td>"
        f"<td>{html.escape(str(row.get('traceContainsFrontierReader', '-')))}</td>"
        f"<td><code>{html.escape(str(row.get('lowByteHandlerVaHex') or '-'))}</code> {html.escape(str(row.get('lowByteHandlerSection') or '-'))}</td>"
        "</tr>"
        for row in summary["tableRows"]
    )
    wrapper_rows = "\n".join(
        "<tr>"
        f"<td>{row.get('step')}</td>"
        f"<td><code>{html.escape(str(row.get('vaHex')))}</code></td>"
        f"<td><code>{html.escape(str(row.get('valueHex')))}</code></td>"
        f"<td><code>{html.escape(str(row.get('opcodeHex')))}</code></td>"
        f"<td><code>{html.escape(str(row.get('handlerVaHex')))}</code></td>"
        f"<td>{html.escape(str(row.get('stopReason') or '-'))}</td>"
        "</tr>"
        for row in summary["wrapperTrace"]
    )
    frontier_rows = "\n".join(
        "<tr>"
        f"<td>{row.get('step')}</td>"
        f"<td><code>{html.escape(str(row.get('vaHex')))}</code></td>"
        f"<td><code>{html.escape(str(row.get('valueHex')))}</code></td>"
        f"<td><code>{html.escape(str(row.get('opcodeHex')))}</code></td>"
        f"<td><code>{html.escape(str(row.get('handlerVaHex')))}</code></td>"
        f"<td><code>{html.escape(str(row.get('branchTargetHex') or '-'))}</code></td>"
        f"<td><code>{html.escape(str(row.get('fallthroughVaHex') or '-'))}</code></td>"
        f"<td>{html.escape(str(row.get('stopReason') or '-'))}</td>"
        "</tr>"
        for row in summary["frontierTrace"]
    )
    proofs = "".join(f"<li>{html.escape(item)}</li>" for item in summary["remainingProofs"])
    missing_items = "".join(
        f"<li>{html.escape(item)}</li>" for item in summary.get("missingEvidence") or []
    )
    return "\n".join([
        "<!doctype html><meta charset=\"utf-8\"><title>Save Selector Leaf Table Context</title>",
        "<style>body{font-family:system-ui,sans-serif;background:#111;color:#eee;max-width:1100px;margin:24px auto}table{border-collapse:collapse;width:100%;margin:18px 0}td,th{border:1px solid #444;padding:6px 8px;vertical-align:top}code{color:#9bd4ff}</style>",
        "<h1>Save Selector Leaf Table Context</h1>",
        f"<p>selector <code>{summary['selector']}</code> root <code>{summary['rootHex']}</code>; root table pointer <code>{summary['rootTablePointerHex']}</code>; table window <code>{summary['tableWindowHex']}</code>.</p>",
        f"<p>wrapper leaf <code>{summary['wrapperLeafHex']}</code>; frontier leaf <code>{summary['frontierLeafHex']}</code>; frontier reader <code>{summary['frontierReaderHex']}</code>.</p>",
        f"<p>leaf refs in window: {summary['leafRefCount']}; frontier leaf ref <code>{summary['frontierLeafRefVaHex']}</code>; frontier leaf is direct root table entry: {summary['frontierLeafRefIsDirectRootTableEntry']}; runtime selection proven: {summary['runtimeSelectionProven']}; proofFound={summary['proofFound']}; failedLeafTableGates=<code>{html.escape(','.join(summary.get('failedLeafTableGateIds') or []) or '-')}</code>; missingEvidenceCount={len(summary.get('missingEvidence') or [])}; promotion status: {html.escape(summary['promotionStatus'])}</p>",
        f"<p>{html.escape(summary['conclusion'])}</p>",
        f"<p><b>Evidence refs:</b> {summary['evidenceRefCount']}.</p>",
        f"<h2>Missing Evidence</h2><ul>{missing_items}</ul>",
        f"<h2>Evidence Refs</h2><ul>{evidence_refs}</ul>",
        "<h2>Table Window</h2><table><thead><tr><th>va</th><th>value</th><th>kind</th><th>leaf maps</th><th>reader?</th><th>handler</th></tr></thead><tbody>",
        table_rows,
        "</tbody></table>",
        "<h2>Wrapper Trace</h2><table><thead><tr><th>step</th><th>va</th><th>value</th><th>opcode</th><th>handler</th><th>stop</th></tr></thead><tbody>",
        wrapper_rows,
        "</tbody></table>",
        "<h2>Frontier Trace</h2><table><thead><tr><th>step</th><th>va</th><th>value</th><th>opcode</th><th>handler</th><th>branch target</th><th>fallthrough</th><th>stop</th></tr></thead><tbody>",
        frontier_rows,
        "</tbody></table>",
        f"<h2>Remaining Proofs</h2><ul>{proofs}</ul>",
    ])


def write_outputs(summary: dict, out_dir: Path, html_out: Path | None = None) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)
    json_out = out_dir / "save_selector_leaf_table_context.json"
    json_out.write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    if html_out is not None:
        html_out.parent.mkdir(parents=True, exist_ok=True)
        html_out.write_text(html_page(summary), encoding="utf-8")
    return json_out


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--exe", type=Path, default=ROOT / "Hwanse2.exe")
    parser.add_argument("--selector-refs", type=Path, default=OUT / "save_scene_selector_references.json")
    parser.add_argument("--leaf-streams", type=Path, default=OUT / "save_selector_leaf_streams.json")
    parser.add_argument("--stream-traces", type=Path, default=OUT / "save_selector_stream_traces.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    parser.add_argument("--html-out", type=Path)
    args = parser.parse_args()
    summary = build_summary(
        args.exe.read_bytes(),
        load_json(args.selector_refs, []),
        load_json(args.leaf_streams, []),
        load_json(args.stream_traces, []),
    )
    json_out = write_outputs(summary, args.out_dir, args.html_out)
    print(f"wrote save selector leaf table context -> {json_out}")


if __name__ == "__main__":
    main()