#!/usr/bin/env python3
"""Summarize dispatch-table entries that reach primaryBranchState writers."""
from __future__ import annotations

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

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

from probe_exe_scene_tables import read_sections, va_to_offset
from summarize_save_selector_branch_state_writers import build_summary as build_writer_summary


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
EVENT_HANDLER_TABLE_VA = 0x0047F1D8
EVENT_HANDLER_COUNT = 64
EVENT_DISPATCHER_VA = 0x0041B687
EVENT_DISPATCH_CALL_VA = 0x0041B6CD
KNOWN_FUNCTION_STARTS = {
    "0x0041dccc..0x0041df2d": 0x0041DC33,
    "0x0041e0f2..0x0041e1f6": 0x0041E03E,
    "0x0041e390..0x0041e431": 0x0041E30F,
    "0x0041fb36..0x0041fb57": 0x0041F96D,
}


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


def section_for_offset(sections: list[dict], offset: int) -> dict | None:
    for section in sections:
        if section["raw"] <= offset < section["raw"] + section["raw_size"]:
            return section
    return None


def section_for_va(sections: list[dict], va: int) -> str | None:
    for section in sections:
        if section["va"] <= va < section["va"] + section["raw_size"]:
            return section["name"]
    return 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 normalize_hex_bytes(text: str) -> str:
    return bytes.fromhex(text).hex(" ")


def bytes_at_va(exe: bytes, sections: list[dict], va: int, length: int) -> str | None:
    offset = va_to_offset(sections, va)
    if offset is None:
        return None
    return exe[offset:offset + length].hex(" ")


def byte_check(exe: bytes, sections: list[dict], va: int, expected: str, label: str) -> dict:
    expected_norm = normalize_hex_bytes(expected)
    actual = bytes_at_va(exe, sections, va, len(bytes.fromhex(expected))) or ""
    return {
        "label": label,
        "vaHex": hex32(va),
        "expectedBytes": expected_norm,
        "actualBytes": actual,
        "matches": actual == expected_norm,
    }


def dispatcher_byte_checks(exe: bytes, sections: list[dict]) -> list[dict]:
    return [
        byte_check(exe, sections, 0x0041B687, "83 3d 80 b2 55 00 00", "loop while dispatch stop flag 0x0055b280 is zero"),
        byte_check(exe, sections, 0x0041B697, "8b 80 b0 00 00 00", "load event/object stream pointer from context+0xb0"),
        byte_check(exe, sections, 0x0041B69F, "8a 08", "read first stream byte"),
        byte_check(exe, sections, 0x0041B6A1, "83 f9 40", "route byte 0x40 to table dispatch"),
        byte_check(exe, sections, 0x0041B6CA, "8a 48 01", "read table index from stream+1"),
        byte_check(exe, sections, 0x0041B6CD, "ff 14 8d d8 f1 47 00", "call dword [index*4 + 0x0047f1d8]"),
    ]


def relative_call_refs(exe: bytes, sections: list[dict], target_va: int) -> list[dict]:
    refs = []
    text = next(section for section in sections if section["name"] == ".text")
    start = text["raw"]
    end = text["raw"] + text["raw_size"]
    raw = exe[start:end]
    for index in range(0, len(raw) - 4):
        if raw[index] != 0xE8:
            continue
        rel = struct.unpack_from("<i", raw, index + 1)[0]
        call_va = text["va"] + index
        resolved = call_va + 5 + rel
        if resolved == target_va:
            refs.append({
                "callVaHex": hex32(call_va),
                "targetVaHex": hex32(target_va),
            })
    return refs


def pointer_refs(exe: bytes, sections: list[dict], target_va: int) -> list[dict]:
    refs = []
    needle = struct.pack("<I", target_va)
    search = 0
    while True:
        hit = exe.find(needle, search)
        if hit < 0:
            break
        search = hit + 1
        section = section_for_offset(sections, hit)
        if section is None:
            continue
        ref_va = section["va"] + hit - section["raw"]
        refs.append({
            "section": section["name"],
            "refVaHex": hex32(ref_va),
            "fileOffsetHex": f"0x{hit:06x}",
        })
    return refs


def event_handler_entries(exe: bytes, sections: list[dict]) -> list[dict]:
    entries = []
    for opcode in range(EVENT_HANDLER_COUNT):
        entry_va = EVENT_HANDLER_TABLE_VA + opcode * 4
        handler_va = dword_at(exe, sections, entry_va)
        entries.append({
            "opcode": opcode,
            "opcodeHex": f"0x{opcode:02x}",
            "entryVaHex": hex32(entry_va),
            "handlerVaHex": hex32(handler_va) if handler_va is not None else None,
            "handlerSection": section_for_va(sections, handler_va) if handler_va is not None else None,
        })
    return entries


def build_summary(exe: bytes, writer_summary: dict | None = None) -> dict:
    sections = read_sections(exe)
    writer_summary = writer_summary or build_writer_summary(exe)
    table_entries = event_handler_entries(exe, sections)
    dispatch_checks = dispatcher_byte_checks(exe, sections)
    entries_by_handler = {entry.get("handlerVaHex"): entry for entry in table_entries}
    clusters = []
    for cluster in writer_summary.get("clusters") or []:
        function_start = KNOWN_FUNCTION_STARTS.get(cluster["rangeHex"])
        function_start_hex = hex32(function_start) if function_start is not None else None
        table_entry = entries_by_handler.get(function_start_hex)
        calls = relative_call_refs(exe, sections, function_start) if function_start is not None else []
        refs = pointer_refs(exe, sections, function_start) if function_start is not None else []
        clusters.append({
            "label": cluster["label"],
            "writerRangeHex": cluster["rangeHex"],
            "functionStartHex": function_start_hex,
            "eventHandlerTable": {
                "tableVaHex": hex32(EVENT_HANDLER_TABLE_VA),
                "entryVaHex": table_entry.get("entryVaHex") if table_entry else None,
                "opcodeHex": table_entry.get("opcodeHex") if table_entry else None,
                "handlerSection": table_entry.get("handlerSection") if table_entry else None,
            },
            "directRelativeCallRefs": calls,
            "directRelativeCallRefCount": len(calls),
            "pointerRefs": refs,
            "pointerRefCount": len(refs),
            "writeValueCounts": cluster.get("writeValueCounts") or {},
            "sourceGlobals": cluster.get("sourceGlobals") or [],
            "interpretation": (
                "Dispatched through the event/object handler table, not by direct relative calls."
                if table_entry and not calls
                else "Direct call evidence exists; inspect call sites before treating this as only data-driven."
            ),
        })
    return {
        "scope": "dispatch paths for primaryBranchState writer clusters",
        "eventHandlerTable": {
            "tableVaHex": hex32(EVENT_HANDLER_TABLE_VA),
            "entryCount": EVENT_HANDLER_COUNT,
            "dispatcherVaHex": hex32(EVENT_DISPATCHER_VA),
            "dispatchCallVaHex": hex32(EVENT_DISPATCH_CALL_VA),
            "dispatcherVerified": all(row["matches"] for row in dispatch_checks),
            "dispatcherByteChecks": dispatch_checks,
            "meaning": (
                "This .data table contains event/object VM handler function pointers. "
                "Dispatcher 0x0041b687 reads byte 0 from context+0xb0; when it is 0x40, "
                "it reads byte 1 as the index and calls dword [index*4 + 0x0047f1d8]. "
                "The branch-state writer routines are registered here by index, so their execution is data-driven."
            ),
        },
        "conclusion": (
            "The direct primaryBranchState writer clusters are table-dispatched event/object handlers. "
            "No direct relative call sites were found for the current writer entry points, so the next route-progress proof "
            "must identify the event/object bytecode opcode that runs before save-selector root 2:0."
        ),
        "clusters": clusters,
        "tableEntries": table_entries,
    }


def format_counts(counts: dict) -> str:
    return ", ".join(f"{key}:{value}" for key, value in sorted(counts.items())) or "-"


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Branch State Dispatch",
        "",
        "Dispatch evidence for direct `primaryBranchState` writer clusters.",
        "",
        f"- event/object handler table: `{summary['eventHandlerTable']['tableVaHex']}`",
        f"- dispatcher: `{summary['eventHandlerTable']['dispatcherVaHex']}` / call `{summary['eventHandlerTable']['dispatchCallVaHex']}`",
        f"- dispatcher byte checks verified: {summary['eventHandlerTable']['dispatcherVerified']}",
        f"- scanned entries: {summary['eventHandlerTable']['entryCount']}",
        "",
        summary["conclusion"],
        "",
        "| cluster | function | table entry | opcode/index | direct calls | writes | interpretation |",
        "| --- | --- | --- | --- | ---: | --- | --- |",
    ]
    for cluster in summary["clusters"]:
        table = cluster["eventHandlerTable"]
        lines.append(
            f"| {cluster['label']} | `{cluster['functionStartHex']}` | "
            f"`{table.get('entryVaHex') or '-'}` | `{table.get('opcodeHex') or '-'}` | "
            f"{cluster['directRelativeCallRefCount']} | {format_counts(cluster['writeValueCounts'])} | "
            f"{cluster['interpretation']} |"
        )
    lines.extend(["", "## Pointer References", ""])
    for cluster in summary["clusters"]:
        refs = ", ".join(f"{ref['section']} `{ref['refVaHex']}`" for ref in cluster["pointerRefs"]) or "-"
        lines.append(f"- `{cluster['functionStartHex']}` {cluster['label']}: {refs}")
    lines.extend(["", "## Dispatcher Byte Checks", ""])
    for check in summary["eventHandlerTable"]["dispatcherByteChecks"]:
        lines.append(
            f"- `{check['vaHex']}` {check['label']}: {check['matches']} "
            f"(`{check['actualBytes']}`)"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    rows = []
    for cluster in summary["clusters"]:
        table = cluster["eventHandlerTable"]
        rows.append(
            "<tr>"
            f"<td>{html.escape(cluster['label'])}</td>"
            f"<td><code>{html.escape(cluster['functionStartHex'] or '-')}</code></td>"
            f"<td><code>{html.escape(table.get('entryVaHex') or '-')}</code></td>"
            f"<td><code>{html.escape(table.get('opcodeHex') or '-')}</code></td>"
            f"<td>{cluster['directRelativeCallRefCount']}</td>"
            f"<td>{html.escape(format_counts(cluster['writeValueCounts']))}</td>"
            f"<td>{html.escape(cluster['interpretation'])}</td>"
            "</tr>"
        )
    ref_rows = []
    for cluster in summary["clusters"]:
        refs = ", ".join(f"{ref['section']} <code>{ref['refVaHex']}</code>" for ref in cluster["pointerRefs"]) or "-"
        ref_rows.append(
            "<tr>"
            f"<td><code>{html.escape(cluster['functionStartHex'] or '-')}</code></td>"
            f"<td>{html.escape(cluster['label'])}</td>"
            f"<td>{refs}</td>"
            "</tr>"
        )
    check_rows = []
    for check in summary["eventHandlerTable"]["dispatcherByteChecks"]:
        check_rows.append(
            "<tr>"
            f"<td><code>{html.escape(check['vaHex'])}</code></td>"
            f"<td>{html.escape(check['label'])}</td>"
            f"<td>{check['matches']}</td>"
            f"<td><code>{html.escape(check['actualBytes'])}</code></td>"
            "</tr>"
        )
    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 Branch State Dispatch</title>",
        "  <style>",
        "    body { margin: 24px; background: #101010; color: #eee; font: 14px system-ui, sans-serif; }",
        "    table { border-collapse: collapse; width: 100%; margin-bottom: 24px; }",
        "    th, td { border: 1px solid #333; padding: 6px 8px; vertical-align: top; }",
        "    th { background: #1d1d1d; }",
        "    code { color: #f5d76e; }",
        "  </style>",
        "</head>",
        "<body>",
        "  <h1>Save Selector Branch State Dispatch</h1>",
        f"  <p>Event/object handler table: <code>{html.escape(summary['eventHandlerTable']['tableVaHex'])}</code>.</p>",
        f"  <p>Dispatcher: <code>{html.escape(summary['eventHandlerTable']['dispatcherVaHex'])}</code>, call <code>{html.escape(summary['eventHandlerTable']['dispatchCallVaHex'])}</code>; dispatcher byte checks verified: {summary['eventHandlerTable']['dispatcherVerified']}.</p>",
        f"  <p>{html.escape(summary['conclusion'])}</p>",
        "  <table><thead><tr><th>cluster</th><th>function</th><th>table entry</th><th>opcode/index</th><th>direct calls</th><th>writes</th><th>interpretation</th></tr></thead>",
        f"  <tbody>{''.join(rows)}</tbody></table>",
        "  <h2>Pointer References</h2>",
        "  <table><thead><tr><th>function</th><th>cluster</th><th>refs</th></tr></thead>",
        f"  <tbody>{''.join(ref_rows)}</tbody></table>",
        "  <h2>Dispatcher Byte Checks</h2>",
        "  <table><thead><tr><th>VA</th><th>check</th><th>matches</th><th>actual bytes</th></tr></thead>",
        f"  <tbody>{''.join(check_rows)}</tbody></table>",
        "</body>",
        "</html>",
        "",
    ])


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


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--exe", type=Path, default=ROOT / "Hwanse2.exe")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    summary = build_summary(args.exe.read_bytes())
    write_outputs(summary, args.out_dir)
    print(f"wrote save selector branch state dispatch -> {args.out_dir / 'save_selector_branch_state_dispatch.html'}")


if __name__ == "__main__":
    main()