#!/usr/bin/env python3
"""Summarize code paths that assign the context+0xa8 selection-buffer base pointer."""
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 offset_to_va, read_sections


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
GATE_OFFSETS = [0x20, 0xE8, 0xEA]
FAILED_SELECTION_BUFFER_BASE_GATE_IDS = [
    "fixed-selection-buffer-base",
    "static-gate-address-ref",
    "runtime-pointer-mode-proof",
]
SELECTION_BUFFER_BASE_MISSING_EVIDENCE = [
    "fixed context+0xa8 base for the current route path",
    "direct static reference to the 0xe8/0xea gate addresses",
    "runtime pointer-mode proof selecting the correct global/save/object base before the gate",
]
SELECTION_BUFFER_BASE_EVIDENCE_REFS = [
    {
        "path": "Hwanse2.exe",
        "description": "context+0xa8 assignment scan and static gate-address reference checks",
    },
]


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


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


def pattern_hits(exe: bytes, sections: list[dict], pattern: bytes) -> list[int]:
    rows = []
    offset = 0
    while True:
        hit = exe.find(pattern, offset)
        if hit < 0:
            return rows
        va = offset_to_va(sections, hit)
        if va is not None:
            rows.append(va)
        offset = hit + 1


def dword_pattern_ref_count(exe: bytes, sections: list[dict], target: int) -> tuple[int, list[str]]:
    pattern = struct.pack("<I", target)
    refs = pattern_hits(exe, sections, pattern)
    return len(refs), [hex32(ref) for ref in refs[:12]]


def immediate_assignments(exe: bytes, sections: list[dict]) -> list[dict]:
    pattern = bytes.fromhex("c780a8000000")
    rows = []
    search = 0
    while True:
        hit = exe.find(pattern, search)
        if hit < 0:
            break
        va = offset_to_va(sections, hit)
        if va is not None and hit + len(pattern) + 4 <= len(exe):
            base = struct.unpack_from("<I", exe, hit + len(pattern))[0]
            rows.append({
                "va": va,
                "vaHex": hex32(va),
                "base": base,
                "baseHex": hex32(base),
                "source": (
                    "global selection buffer"
                    if base == 0x0059E310
                    else "save/runtime block"
                    if base == 0x004576D8
                    else "immediate base"
                ),
                "instruction": "mov dword [eax+0xa8], imm32",
            })
        search = hit + 1
    return rows


def register_assignments(exe: bytes, sections: list[dict]) -> list[dict]:
    pattern = bytes.fromhex("8981a8000000")
    known = {
        0x0040575A: {
            "source": "saved pointer table 0x0059dd70[stream+2]",
            "baseExpression": "dword[0x0059dd70 + stream[2]*4]",
        },
        0x0040643D: {
            "source": "save block slot selected by stream+1",
            "baseExpression": "0x00457750 + stream[1] * 0xd8",
        },
        0x00406480: {
            "source": "runtime object pointer table selected by context+0xf2",
            "baseExpression": "dword[0x0059db30 + context[0xf2]*4]",
        },
        0x004064A0: {
            "source": "runtime object pointer table selected by stream+2",
            "baseExpression": "dword[0x0059db30 + stream[2]*4]",
        },
        0x0041D8FF: {
            "source": "save block slot selected by object field +3",
            "baseExpression": "0x00457750 + object[3] * 0xd8",
        },
        0x0041D91F: {
            "source": "runtime object pointer table selected by context+0xf2",
            "baseExpression": "dword[0x0059db30 + context[0xf2]*4]",
        },
    }
    return [
        {
            "va": va,
            "vaHex": hex32(va),
            "instruction": "mov dword [ecx+0xa8], eax",
            **known.get(va, {"source": "register-computed base", "baseExpression": "eax"}),
        }
        for va in pattern_hits(exe, sections, pattern)
    ]


def gate_address_rows(exe: bytes, sections: list[dict]) -> list[dict]:
    bases = [
        {
            "baseHex": "0x0059e310",
            "base": 0x0059E310,
            "source": "global selection buffer immediate assignment",
            "static": True,
        },
        {
            "baseHex": "0x004576d8",
            "base": 0x004576D8,
            "source": "save/runtime block immediate assignment",
            "static": True,
        },
        {
            "baseHex": "0x00457750+n*0xd8",
            "base": 0x00457750,
            "source": "indexed save block slot base",
            "static": False,
        },
        {
            "baseHex": "dword[0x0059db30+i*4]",
            "base": None,
            "source": "runtime object pointer table base",
            "static": False,
        },
        {
            "baseHex": "dword[0x0059dd70+i*4]",
            "base": None,
            "source": "saved pointer table base",
            "static": False,
        },
    ]
    rows = []
    for base in bases:
        for offset in GATE_OFFSETS:
            address = base["base"] + offset if isinstance(base.get("base"), int) else None
            count = 0
            refs: list[str] = []
            if address is not None:
                count, refs = dword_pattern_ref_count(exe, sections, address)
            rows.append({
                "baseHex": base["baseHex"],
                "baseSource": base["source"],
                "offsetHex": hex8(offset),
                "addressHex": hex32(address) if address is not None else None,
                "directDwordRefCount": count,
                "directDwordRefs": refs,
                "staticAddress": base["static"],
            })
    return rows


def build_summary(exe: bytes) -> dict:
    sections = read_sections(exe)
    immediates = immediate_assignments(exe, sections)
    registers = register_assignments(exe, sections)
    address_rows = gate_address_rows(exe, sections)
    known_static_e8_ea_refs = [
        row for row in address_rows
        if row["offsetHex"] in {"0xe8", "0xea"}
        and row["staticAddress"]
        and row["directDwordRefCount"] > 0
    ]
    conclusion = (
        "The selection-buffer pointer at context+0xa8 is not a single fixed array. Static handler paths can "
        "point it at 0x0059e310 or 0x004576d8, while other paths compute a save-slot base or load an object "
        "pointer from 0x0059db30/0x0059dd70. The gate offsets 0xe8/0xea therefore need runtime pointer-mode "
        "evidence before they can prove fallthrough. Exact dword refs to the static 0x0059e310+0xe8/0xea and "
        "0x004576d8+0xe8/0xea addresses are absent, so the producer is still unresolved."
    )
    return {
        "contextFieldHex": "0x000000a8",
        "immediateAssignmentCount": len(immediates),
        "registerAssignmentCount": len(registers),
        "immediateAssignments": immediates,
        "registerAssignments": registers,
        "gateAddressRows": address_rows,
        "knownStaticGateOffsetDirectRefCount": len(known_static_e8_ea_refs),
        "runtimePointerModeStillRequired": True,
        "proofFound": False,
        "selectionBufferBaseProofFound": False,
        "failedSelectionBufferBaseGateIds": FAILED_SELECTION_BUFFER_BASE_GATE_IDS,
        "missingEvidence": SELECTION_BUFFER_BASE_MISSING_EVIDENCE,
        "evidenceRefs": SELECTION_BUFFER_BASE_EVIDENCE_REFS,
        "evidenceRefCount": len(SELECTION_BUFFER_BASE_EVIDENCE_REFS),
        "promotionStatus": "blocked",
        "conclusion": conclusion,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Selection Buffer Bases",
        "",
        f"- context field: `{summary['contextFieldHex']}`",
        f"- immediate assignments: {summary['immediateAssignmentCount']}",
        f"- register/computed assignments: {summary['registerAssignmentCount']}",
        f"- static gate-offset direct refs: {summary['knownStaticGateOffsetDirectRefCount']}",
        f"- runtime pointer mode still required: {summary['runtimePointerModeStillRequired']}",
        f"- proofFound: `{summary['proofFound']}`",
        f"- selectionBufferBaseProofFound: `{summary['selectionBufferBaseProofFound']}`",
        f"- promotion status: `{summary['promotionStatus']}`",
        "",
        summary["conclusion"],
        "",
        "## Failed Gates",
        "",
    ]
    lines.extend(f"- `{item}`" for item in summary["failedSelectionBufferBaseGateIds"])
    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([
        "",
        "## Immediate Assignments",
        "",
        "| va | base | source | instruction |",
        "| --- | --- | --- | --- |",
    ])
    for row in summary["immediateAssignments"]:
        lines.append(f"| `{row['vaHex']}` | `{row['baseHex']}` | {row['source']} | `{row['instruction']}` |")
    lines.extend([
        "",
        "## Register Assignments",
        "",
        "| va | source | base expression | instruction |",
        "| --- | --- | --- | --- |",
    ])
    for row in summary["registerAssignments"]:
        lines.append(
            f"| `{row['vaHex']}` | {row['source']} | `{row['baseExpression']}` | `{row['instruction']}` |"
        )
    lines.extend([
        "",
        "## Gate Offset Address Checks",
        "",
        "| base | source | offset | address | direct dword refs | sample refs |",
        "| --- | --- | --- | --- | ---: | --- |",
    ])
    for row in summary["gateAddressRows"]:
        lines.append(
            f"| `{row['baseHex']}` | {row['baseSource']} | `{row['offsetHex']}` | "
            f"`{row['addressHex'] or '-'}` | {row['directDwordRefCount']} | "
            f"{', '.join(row['directDwordRefs']) or '-'} |"
        )
    lines.append("")
    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["failedSelectionBufferBaseGateIds"]
    )
    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"]
    )
    immediate_rows = []
    for row in summary["immediateAssignments"]:
        immediate_rows.append(
            "<tr>"
            f"<td><code>{html.escape(row['vaHex'])}</code></td>"
            f"<td><code>{html.escape(row['baseHex'])}</code></td>"
            f"<td>{html.escape(row['source'])}</td>"
            f"<td><code>{html.escape(row['instruction'])}</code></td>"
            "</tr>"
        )
    register_rows = []
    for row in summary["registerAssignments"]:
        register_rows.append(
            "<tr>"
            f"<td><code>{html.escape(row['vaHex'])}</code></td>"
            f"<td>{html.escape(row['source'])}</td>"
            f"<td><code>{html.escape(row['baseExpression'])}</code></td>"
            f"<td><code>{html.escape(row['instruction'])}</code></td>"
            "</tr>"
        )
    address_rows = []
    for row in summary["gateAddressRows"]:
        address_rows.append(
            "<tr>"
            f"<td><code>{html.escape(row['baseHex'])}</code></td>"
            f"<td>{html.escape(row['baseSource'])}</td>"
            f"<td><code>{html.escape(row['offsetHex'])}</code></td>"
            f"<td><code>{html.escape(row['addressHex'] or '-')}</code></td>"
            f"<td>{row['directDwordRefCount']}</td>"
            f"<td>{html.escape(', '.join(row['directDwordRefs']) or '-')}</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 Selection Buffer Bases</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 Selection Buffer Bases</h1>",
        f"  <p>context field <code>{html.escape(summary['contextFieldHex'])}</code>; immediate assignments {summary['immediateAssignmentCount']}; register/computed assignments {summary['registerAssignmentCount']}; static gate-offset direct refs {summary['knownStaticGateOffsetDirectRefCount']}; runtime pointer mode still required: {summary['runtimePointerModeStillRequired']}; proofFound <code>{summary['proofFound']}</code>; promotion status: {html.escape(summary['promotionStatus'])}.</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>Immediate Assignments</h2>",
        "  <table><thead><tr><th>va</th><th>base</th><th>source</th><th>instruction</th></tr></thead><tbody>",
        *immediate_rows,
        "  </tbody></table>",
        "  <h2>Register Assignments</h2>",
        "  <table><thead><tr><th>va</th><th>source</th><th>base expression</th><th>instruction</th></tr></thead><tbody>",
        *register_rows,
        "  </tbody></table>",
        "  <h2>Gate Offset Address Checks</h2>",
        "  <table><thead><tr><th>base</th><th>source</th><th>offset</th><th>address</th><th>direct dword refs</th><th>sample refs</th></tr></thead><tbody>",
        *address_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_selection_buffer_bases.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (out_dir / "save_selector_selection_buffer_bases.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 selection buffer bases -> {args.out_dir / 'save_selector_selection_buffer_bases.html'}")


if __name__ == "__main__":
    main()