#!/usr/bin/env python3
"""Summarize runtime state tables used by save-selector stream branches."""
from __future__ import annotations

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

from probe_exe_pointer_refs import find_value_refs
from probe_exe_scene_tables import offset_to_va, read_sections


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
STATE_TABLES = [
    {
        "name": "primaryBranchState",
        "baseVa": 0x0059E370,
        "slots": 12,
        "usedWhen": "opcode 0x11 stream+1 == 0",
    },
    {
        "name": "secondaryBranchState",
        "baseVa": 0x0059E360,
        "slots": 12,
        "usedWhen": "opcode 0x11 stream+1 != 0",
    },
]


def classify_ref(exe: bytes, ref: dict) -> dict:
    offset = ref["fileOffset"]
    value = ref["value"]
    before = exe[max(0, offset - 4): offset]
    after = exe[offset + 4: offset + 12]
    item = {
        "section": ref["section"],
        "fileOffset": offset,
        "fileOffsetHex": f"0x{offset:06x}",
        "refVa": ref["refVa"],
        "refVaHex": f"0x{ref['refVa']:08x}",
        "valueHex": f"0x{value:08x}",
        "kind": "unknown",
    }
    # c6 80 <disp32> <imm8> = mov byte ptr [eax + disp32], imm8
    if offset >= 2 and exe[offset - 2: offset] == b"\xc6\x80" and after:
        item["kind"] = "indexedWriteImmediate"
        item["writeValue"] = after[0]
        item["writeValueHex"] = f"0x{after[0]:02x}"
    # 8a 81 <disp32> = mov al, byte ptr [ecx + disp32]
    elif offset >= 2 and exe[offset - 2: offset] == b"\x8a\x81":
        item["kind"] = "indexedReadByte"
    # b8/b9 <imm32> are base address constants passed to branch handlers.
    elif before[-1:] in {b"\xb8", b"\xb9"}:
        item["kind"] = "baseAddressImmediate"
    return item


def build_summary(exe: bytes) -> dict:
    sections = read_sections(exe)
    tables = []
    for table in STATE_TABLES:
        refs = [
            classify_ref(exe, ref)
            for ref in find_value_refs(exe, sections, table["baseVa"], {".text"})
        ]
        by_kind: dict[str, int] = {}
        write_values: dict[str, int] = {}
        for ref in refs:
            by_kind[ref["kind"]] = by_kind.get(ref["kind"], 0) + 1
            if ref.get("writeValueHex"):
                write_values[ref["writeValueHex"]] = write_values.get(ref["writeValueHex"], 0) + 1
        tables.append({
            **table,
            "baseVaHex": f"0x{table['baseVa']:08x}",
            "rangeHex": f"0x{table['baseVa']:08x}..0x{table['baseVa'] + table['slots'] - 1:08x}",
            "refs": refs,
            "refKindCounts": by_kind,
            "writeValueCounts": write_values,
        })
    return {
        "scope": "runtime branch state tables referenced by save-selector opcode 0x11",
        "tables": tables,
        "opcode11": {
            "handlerVaHex": "0x0040b4e6",
            "condition": (
                "read table[(context+0xa8)[stream+2]]; value 1 falls through +8, "
                "any other value jumps to dword [stream+4]"
            ),
        },
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Branch State",
        "",
        "Runtime state tables used by save-selector opcode `0x11`.",
        "",
        f"Opcode 0x11 handler: `{summary['opcode11']['handlerVaHex']}`.",
        summary["opcode11"]["condition"],
        "",
        "| table | range | used when | refs | ref kinds | write values |",
        "| --- | --- | --- | ---: | --- | --- |",
    ]
    for table in summary["tables"]:
        kinds = ", ".join(f"{key}:{value}" for key, value in sorted(table["refKindCounts"].items())) or "-"
        writes = ", ".join(f"{key}:{value}" for key, value in sorted(table["writeValueCounts"].items())) or "-"
        lines.append(
            f"| {table['name']} | `{table['rangeHex']}` | {table['usedWhen']} | "
            f"{len(table['refs'])} | {kinds} | {writes} |"
        )
    lines.extend(["", "## References", ""])
    for table in summary["tables"]:
        lines.extend([
            f"### {table['name']} {table['rangeHex']}",
            "",
            "| ref | kind | write value |",
            "| --- | --- | --- |",
        ])
        for ref in table["refs"]:
            lines.append(
                f"| `{ref['refVaHex']}` | {ref['kind']} | `{ref.get('writeValueHex') or '-'}` |"
            )
        lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    rows = []
    for table in summary["tables"]:
        kinds = ", ".join(f"{key}:{value}" for key, value in sorted(table["refKindCounts"].items())) or "-"
        writes = ", ".join(f"{key}:{value}" for key, value in sorted(table["writeValueCounts"].items())) or "-"
        rows.append(
            "<tr>"
            f"<td>{html.escape(table['name'])}</td>"
            f"<td><code>{html.escape(table['rangeHex'])}</code></td>"
            f"<td>{html.escape(table['usedWhen'])}</td>"
            f"<td>{len(table['refs'])}</td>"
            f"<td>{html.escape(kinds)}</td>"
            f"<td>{html.escape(writes)}</td>"
            "</tr>"
        )
    ref_sections = []
    for table in summary["tables"]:
        body = []
        for ref in table["refs"]:
            body.append(
                "<tr>"
                f"<td><code>{html.escape(ref['refVaHex'])}</code></td>"
                f"<td>{html.escape(ref['kind'])}</td>"
                f"<td><code>{html.escape(ref.get('writeValueHex') or '-')}</code></td>"
                "</tr>"
            )
        ref_sections.append(
            f"<h2>{html.escape(table['name'])} {html.escape(table['rangeHex'])}</h2>"
            "<table><thead><tr><th>ref</th><th>kind</th><th>write value</th></tr></thead>"
            f"<tbody>{''.join(body)}</tbody></table>"
        )
    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</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</h1>",
        f"  <p>Opcode 0x11 handler: <code>{html.escape(summary['opcode11']['handlerVaHex'])}</code>.</p>",
        f"  <p>{html.escape(summary['opcode11']['condition'])}</p>",
        "  <table><thead><tr><th>table</th><th>range</th><th>used when</th><th>refs</th><th>ref kinds</th><th>write values</th></tr></thead>",
        f"  <tbody>{''.join(rows)}</tbody></table>",
        "\n".join(ref_sections),
        "</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.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        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 -> {args.out_dir / 'save_selector_branch_state.json'}")


if __name__ == "__main__":
    main()