#!/usr/bin/env python3
"""Document the original EXE save-loader selector path."""
from __future__ import annotations

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

from probe_exe_scene_tables import read_sections, va_to_offset


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


EXPECTED_BYTES = {
    0x004234A1: bytes.fromhex("33c0"),
    0x004234A3: bytes.fromhex("a0da764500"),
    0x004234A8: bytes.fromhex("8b0485352d4400"),
    0x004234AF: bytes.fromhex("33c9"),
    0x004234B1: bytes.fromhex("8a0ddb764500"),
    0x004234B7: bytes.fromhex("8b0488"),
    0x004234BA: bytes.fromhex("a330de5900"),
    0x0040ADB8: bytes.fromhex("a330de5900"),
    0x0040ADFE: bytes.fromhex("a130de5900"),
    0x0040AE06: bytes.fromhex("894140"),
    0x0040AE37: bytes.fromhex("a330de5900"),
    0x0040AE4A: bytes.fromhex("a330de5900"),
}


def bytes_at(exe: bytes, sections: list[dict], va: int, size: int) -> bytes:
    offset = va_to_offset(sections, va)
    if offset is None:
        raise ValueError(f"0x{va:08x} is not in a loaded section")
    return exe[offset : offset + size]


def verify_bytes(exe: bytes, sections: list[dict]) -> list[dict]:
    rows = []
    for va, expected in EXPECTED_BYTES.items():
        actual = bytes_at(exe, sections, va, len(expected))
        if actual != expected:
            raise ValueError(f"0x{va:08x}: expected {expected.hex()}, got {actual.hex()}")
        rows.append({"vaHex": f"0x{va:08x}", "bytes": actual.hex(" ")})
    return rows


def dword_at(exe: bytes, sections: list[dict], va: int) -> int:
    return struct.unpack("<I", bytes_at(exe, sections, va, 4))[0]


def build_summary(exe: bytes) -> dict:
    sections = read_sections(exe)
    selector_instruction_bytes = verify_bytes(exe, sections)
    selector_table = dword_at(exe, sections, 0x00442D35 + 2 * 4)
    selected_pointer = dword_at(exe, sections, selector_table)
    second_level_pointer = dword_at(exe, sections, selected_pointer + 2 * 4)
    selector_leaf = dword_at(exe, sections, second_level_pointer + 6 * 4)
    return {
        "saveReadBlocks": [
            {
                "callVaHex": "0x00423404",
                "targetVaHex": "0x004576d8",
                "saveOffsetHex": "0x0000",
                "sizeHex": "0x0072",
                "description": "first save block, includes scene selector bytes and tile/camera candidates",
            },
            {
                "callVaHex": "0x00423430",
                "targetVaHex": "0x00457750",
                "saveOffsetHex": "0x0072",
                "sizeHex": "0x0288",
                "description": "second save block",
            },
            {
                "callVaHex": "0x0042345c",
                "targetVaHex": "0x0059db60",
                "saveOffsetHex": "0x02fa",
                "sizeHex": "0x0200",
                "description": "third save block",
            },
        ],
        "selectorFields": [
            {
                "saveOffsetHex": "0x0002",
                "runtimeVaHex": "0x004576da",
                "field": "scene selector group",
                "loaderVaHex": "0x004234a3",
            },
            {
                "saveOffsetHex": "0x0003",
                "runtimeVaHex": "0x004576db",
                "field": "scene selector slot",
                "loaderVaHex": "0x004234b1",
            },
            {
                "saveOffsetHex": "0x0004",
                "runtimeVaHex": "0x004576dc",
                "field": "tile/camera X candidate",
                "loaderVaHex": "0x00424a93",
            },
            {
                "saveOffsetHex": "0x0006",
                "runtimeVaHex": "0x004576de",
                "field": "tile/camera Y candidate",
                "loaderVaHex": "0x00424ab8",
            },
        ],
        "selectorTrace": [
            {
                "vaHex": "0x004234a3",
                "operation": "AL = byte [0x004576da]",
                "meaning": "load save selector group from save offset 0x0002",
            },
            {
                "vaHex": "0x004234a8",
                "operation": "EAX = dword [0x00442d35 + AL * 4]",
                "meaning": "select row pointer from group table",
            },
            {
                "vaHex": "0x004234b1",
                "operation": "CL = byte [0x004576db]",
                "meaning": "load save selector slot from save offset 0x0003",
            },
            {
                "vaHex": "0x004234b7",
                "operation": "EAX = dword [EAX + ECX * 4]",
                "meaning": "select scene/progress pointer from row and slot",
            },
            {
                "vaHex": "0x004234ba",
                "operation": "dword [0x0059de30] = EAX",
                "meaning": "store selected scene/progress pointer for later runtime use",
            },
        ],
        "selectedPointerRuntimeUses": [
            {
                "vaHex": "0x0040ad9b",
                "handlerTableVaHex": "0x0044073c",
                "handlerIndex": 7,
                "kind": "selectIndexedPointer",
                "operation": "stream byte + table pointer -> dword [0x0059de30]",
                "meaning": "runtime script opcode selects an indexed pointer into the same selected-pointer global",
                "evidenceVaHex": "0x0040adb8",
            },
            {
                "vaHex": "0x0040adc9",
                "handlerTableVaHex": "0x00440740",
                "handlerIndex": 8,
                "kind": "jumpToSelectedPointer",
                "operation": "if dword [0x0059de30] != 0, save current stream and set context+0x40 = dword [0x0059de30]",
                "meaning": "selected pointer becomes the active script/event stream pointer",
                "evidenceVaHex": "0x0040adfe",
            },
            {
                "vaHex": "0x0040ae0e",
                "handlerTableVaHex": "0x00440744",
                "handlerIndex": 9,
                "kind": "setSelectedPointer",
                "operation": "mode 0 stores current stream pointer; mode 1 stores dword [stream+4] into 0x0059de30",
                "meaning": "runtime script opcodes update the selected pointer global",
                "evidenceVaHex": "0x0040ae37",
            },
        ],
        "scriptHandlerTable": {
            "startVaHex": "0x00440720",
            "selectedPointerHandlers": [
                {
                    "index": 7,
                    "entryVaHex": "0x0044073c",
                    "handlerVaHex": "0x0040ad9b",
                    "kind": "selectIndexedPointer",
                },
                {
                    "index": 8,
                    "entryVaHex": "0x00440740",
                    "handlerVaHex": "0x0040adc9",
                    "kind": "jumpToSelectedPointer",
                },
                {
                    "index": 9,
                    "entryVaHex": "0x00440744",
                    "handlerVaHex": "0x0040ae0e",
                    "kind": "setSelectedPointer",
                },
            ],
        },
        "selectorInstructionBytes": selector_instruction_bytes,
        "frontierExample": {
            "group": 2,
            "slot": 0,
            "groupTableVaHex": "0x00442d35",
            "rowPointerVaHex": f"0x{selector_table:08x}",
            "selectedPointerHex": f"0x{selected_pointer:08x}",
            "secondLevelPointerHex": f"0x{second_level_pointer:08x}",
            "leafPointerHex": f"0x{selector_leaf:08x}",
            "knownFrontier": "map1_01a -> map2_02d via leaf 0x00542ac0",
        },
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Loader Trace",
        "",
        "Original EXE save-load evidence for the scene selector fields.",
        "",
        "## Save Read Blocks",
        "",
        "| call | target | save offset | size | description |",
        "| --- | --- | --- | --- | --- |",
    ]
    for row in summary["saveReadBlocks"]:
        lines.append(
            f"| {row['callVaHex']} | {row['targetVaHex']} | {row['saveOffsetHex']} | {row['sizeHex']} | {row['description']} |"
        )
    lines.extend([
        "",
        "## Selector Fields",
        "",
        "| save offset | runtime VA | field | evidence VA |",
        "| --- | --- | --- | --- |",
    ])
    for row in summary["selectorFields"]:
        lines.append(f"| {row['saveOffsetHex']} | {row['runtimeVaHex']} | {row['field']} | {row['loaderVaHex']} |")
    lines.extend([
        "",
        "## Selector Trace",
        "",
        "| VA | operation | meaning |",
        "| --- | --- | --- |",
    ])
    for row in summary["selectorTrace"]:
        lines.append(f"| {row['vaHex']} | `{row['operation']}` | {row['meaning']} |")
    example = summary["frontierExample"]
    lines.extend([
        "",
        "## Runtime Uses Of Selected Pointer",
        "",
        f"Handler table starts at `{summary['scriptHandlerTable']['startVaHex']}`.",
        "",
        "| index | table entry | handler VA | kind | operation | meaning | evidence |",
        "| ---: | --- | --- | --- | --- | --- | --- |",
    ])
    for row in summary["selectedPointerRuntimeUses"]:
        lines.append(
            f"| {row['handlerIndex']} | {row['handlerTableVaHex']} | {row['vaHex']} | {row['kind']} | "
            f"`{row['operation']}` | {row['meaning']} | {row['evidenceVaHex']} |"
        )
    lines.extend([
        "",
        "## Frontier Example",
        "",
        f"- group/slot: `{example['group']}:{example['slot']}`",
        f"- group table: `{example['groupTableVaHex']}`",
        f"- row pointer: `{example['rowPointerVaHex']}`",
        f"- selected pointer: `{example['selectedPointerHex']}`",
        f"- second-level pointer: `{example['secondLevelPointerHex']}`",
        f"- sample leaf pointer: `{example['leafPointerHex']}`",
        f"- known frontier: {example['knownFrontier']}",
        "",
    ])
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    def table(headers: list[str], rows: list[list[str]]) -> str:
        head = "".join(f"<th>{html.escape(header)}</th>" for header in headers)
        body = "\n".join(
            "<tr>" + "".join(f"<td>{html.escape(value)}</td>" for value in row) + "</tr>"
            for row in rows
        )
        return f"<table><thead><tr>{head}</tr></thead><tbody>{body}</tbody></table>"

    read_rows = [
        [row["callVaHex"], row["targetVaHex"], row["saveOffsetHex"], row["sizeHex"], row["description"]]
        for row in summary["saveReadBlocks"]
    ]
    field_rows = [
        [row["saveOffsetHex"], row["runtimeVaHex"], row["field"], row["loaderVaHex"]]
        for row in summary["selectorFields"]
    ]
    trace_rows = [
        [row["vaHex"], row["operation"], row["meaning"]]
        for row in summary["selectorTrace"]
    ]
    runtime_rows = [
        [
            str(row["handlerIndex"]),
            row["handlerTableVaHex"],
            row["vaHex"],
            row["kind"],
            row["operation"],
            row["meaning"],
            row["evidenceVaHex"],
        ]
        for row in summary["selectedPointerRuntimeUses"]
    ]
    example = summary["frontierExample"]
    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 Loader Trace</title>",
        "  <style>",
        "    body { margin: 24px; background: #101010; color: #eee; font: 14px system-ui, sans-serif; }",
        "    table { border-collapse: collapse; width: 100%; margin: 12px 0 24px; }",
        "    th, td { border: 1px solid #333; padding: 6px 8px; vertical-align: top; }",
        "    th { background: #1d1d1d; }",
        "    code { color: #d6e8ff; }",
        "  </style>",
        "</head>",
        "<body>",
        "  <h1>Save Loader Trace</h1>",
        "  <p>Original EXE save-load evidence for the scene selector fields.</p>",
        "  <h2>Save Read Blocks</h2>",
        table(["call", "target", "save offset", "size", "description"], read_rows),
        "  <h2>Selector Fields</h2>",
        table(["save offset", "runtime VA", "field", "evidence VA"], field_rows),
        "  <h2>Selector Trace</h2>",
        table(["VA", "operation", "meaning"], trace_rows),
        "  <h2>Runtime Uses Of Selected Pointer</h2>",
        f"  <p>Handler table starts at <code>{summary['scriptHandlerTable']['startVaHex']}</code>.</p>",
        table(["index", "table entry", "handler VA", "kind", "operation", "meaning", "evidence"], runtime_rows),
        "  <h2>Frontier Example</h2>",
        f"  <p><code>{example['group']}:{example['slot']}</code> resolves through <code>{example['groupTableVaHex']}</code> "
        f"to <code>{example['rowPointerVaHex']}</code>, selected pointer <code>{example['selectedPointerHex']}</code>, "
        f"second-level pointer <code>{example['secondLevelPointerHex']}</code>, then sample leaf <code>{example['leafPointerHex']}</code>.</p>",
        f"  <p>{html.escape(example['knownFrontier'])}</p>",
        "</body>",
        "</html>",
        "",
    ])


def write_outputs(summary: dict, out_dir: Path = OUT) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_loader_trace.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 loader trace -> {args.out_dir / 'save_loader_trace.json'}")


if __name__ == "__main__":
    main()