#!/usr/bin/env python3
"""Summarize original EXE save-file path construction and loader API usage."""
from __future__ import annotations

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

from probe_exe_scene_tables import c_string, offset_to_va, read_sections, va_to_offset
from summarize_exe_imports import parse_imports


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

TARGET_STRING_VAS = {
    0x0048B0C0: "savedat wildcard for slot scan",
    0x0048B0E0: "SaveData fallback directory for wildcard scan",
    0x0048B104: "savedat0 base for current slot",
    0x0048B124: "SaveData fallback directory for current slot",
    0x0048B148: "savedat0 base for explicit slot",
    0x0048B168: "SaveData fallback directory for explicit slot",
    0x004A6590: "registry SavePath value writer",
    0x004A65B0: "registry SavePath value reader",
}

INTERESTING_APIS = {
    "CreateFileA",
    "ReadFile",
    "WriteFile",
    "CloseHandle",
    "FindFirstFileA",
    "FindClose",
    "lstrcpyA",
    "lstrcatA",
    "RegOpenKeyExA",
    "RegQueryValueExA",
    "RegSetValueExA",
    "RegCloseKey",
}

EXPECTED_BYTES = {
    0x00423319: bytes.fromhex("558bec83ec24535657"),
    0x00423337: bytes.fromhex("6848b148006800b35500ff1500045a00"),
    0x0042334F: bytes.fromhex("a207b35500"),
    0x00423366: bytes.fromhex("e8d8320000"),
    0x004233CF: bytes.fromhex("e817050000"),
    0x00423404: bytes.fromhex("ff151c045a00"),
    0x00423430: bytes.fromhex("ff151c045a00"),
    0x0042345C: bytes.fromhex("ff151c045a00"),
    0x00423478: bytes.fromhex("ff1568045a00"),
    0x004234A3: bytes.fromhex("a0da764500"),
    0x004234B1: bytes.fromhex("8a0ddb764500"),
    0x004234BA: bytes.fromhex("a330de5900"),
    0x004238EB: bytes.fromhex("558bec83ec0c535657"),
    0x00423917: bytes.fromhex("ff15e8035a00"),
    0x0042393B: bytes.fromhex("ff15e8035a00"),
    0x0042395F: bytes.fromhex("ff15e8035a00"),
    0x004230EB: bytes.fromhex("ff15f8035a00"),
    0x004266FD: bytes.fromhex("ff1558035a00"),
}


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


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


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"{hex_va(va)} is not in a loaded raw section")
    return exe[offset : offset + size]


def verify_expected_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"{hex_va(va)}: expected {expected.hex()}, got {actual.hex()}")
        rows.append({"vaHex": hex_va(va), "bytes": actual.hex(" ")})
    return rows


def find_dword_refs(exe: bytes, sections: list[dict], value: int) -> list[dict]:
    needle = struct.pack("<I", value)
    refs = []
    search = 0
    while True:
        hit = exe.find(needle, search)
        if hit < 0:
            break
        va = offset_to_va(sections, hit)
        section = section_for_offset(sections, hit)
        if va is not None and section is not None:
            refs.append(
                {
                    "vaHex": hex_va(va),
                    "section": section,
                    "fileOffsetHex": f"0x{hit:06x}",
                }
            )
        search = hit + 1
    return refs


def import_lookup(imports: dict) -> dict[str, dict]:
    rows = {}
    for dll in imports.get("imports") or []:
        dll_name = dll.get("dll", "")
        for function in dll.get("functions") or []:
            name = function.get("name")
            if name:
                rows[name] = {**function, "dll": dll_name}
    return rows


def iat_uses(exe: bytes, sections: list[dict], iat_va: int) -> list[dict]:
    text = next(section for section in sections if section["name"] == ".text")
    raw = exe[text["raw"] : text["raw"] + text["raw_size"]]
    needle = struct.pack("<I", iat_va)
    rows = []
    search = 0
    while True:
        hit = raw.find(needle, search)
        if hit < 0:
            break
        opcode_offset = text["raw"] + hit
        kind = "raw-iat-immediate"
        start_offset = opcode_offset
        if hit >= 2 and raw[hit - 2 : hit] == b"\xff\x15":
            kind = "call-indirect"
            start_offset = text["raw"] + hit - 2
        elif hit >= 2 and raw[hit - 2 : hit] == b"\xff\x25":
            kind = "jmp-indirect"
            start_offset = text["raw"] + hit - 2
        start_va = offset_to_va(sections, start_offset)
        rows.append({"vaHex": hex_va(start_va), "kind": kind})
        search = hit + 1
    return rows


def save_path_strings(exe: bytes, sections: list[dict]) -> list[dict]:
    rows = []
    for va, meaning in TARGET_STRING_VAS.items():
        offset = va_to_offset(sections, va)
        if offset is None:
            raise ValueError(f"save string {hex_va(va)} is not mapped")
        text = c_string(exe, offset)
        refs = find_dword_refs(exe, sections, va)
        rows.append(
            {
                "vaHex": hex_va(va),
                "fileOffsetHex": f"0x{offset:06x}",
                "text": text,
                "meaning": meaning,
                "directRefCount": len(refs),
                "directRefs": refs,
            }
        )
    return rows


def api_summary(exe: bytes, sections: list[dict], imports: dict) -> list[dict]:
    lookup = import_lookup(imports)
    rows = []
    for name in sorted(INTERESTING_APIS):
        item = lookup.get(name)
        if not item:
            continue
        uses = iat_uses(exe, sections, item["iatVa"])
        call_sites = [use for use in uses if use["kind"] == "call-indirect"]
        rows.append(
            {
                "dll": item["dll"],
                "name": name,
                "iatVaHex": hex_va(item["iatVa"]),
                "useCount": len(uses),
                "callCount": len(call_sites),
                "callSites": call_sites,
            }
        )
    return rows


def build_summary(exe: bytes, imports: dict | None = None) -> dict:
    sections = read_sections(exe)
    imports = imports if imports is not None else parse_imports(ROOT / "Hwanse2.exe")
    strings = save_path_strings(exe, sections)
    apis = api_summary(exe, sections, imports)
    api_by_name = {row["name"]: row for row in apis}
    read_calls = [
        {"callVaHex": "0x00423404", "targetVaHex": "0x004576d8", "saveOffsetHex": "0x0000", "sizeHex": "0x0072"},
        {"callVaHex": "0x00423430", "targetVaHex": "0x00457750", "saveOffsetHex": "0x0072", "sizeHex": "0x0288"},
        {"callVaHex": "0x0042345c", "targetVaHex": "0x0059db60", "saveOffsetHex": "0x02fa", "sizeHex": "0x0200"},
    ]
    path_builders = [
        {
            "name": "wildcard slot scan path",
            "functionVaHex": "0x00423041",
            "filenameStringVaHex": "0x0048b0c0",
            "directoryStringVaHex": "0x0048b0e0",
            "outputBufferVaHex": "0x0055b420",
            "apiEvidence": "FindFirstFileA at 0x004230eb",
            "purpose": "check whether any savedat?.dat exists under registry SavePath or SaveData\\ fallback",
        },
        {
            "name": "explicit slot path",
            "functionVaHex": "0x00423319",
            "filenameStringVaHex": "0x0048b148",
            "directoryStringVaHex": "0x0048b168",
            "outputBufferVaHex": "0x0055b318",
            "apiEvidence": "lstrcpyA/lstrcatA build path; open helper called at 0x004233cf",
            "purpose": "build SaveData\\savedatN.dat for the slot argument and feed the loader/saver",
        },
    ]
    loader = {
        "functionVaHex": "0x00423319",
        "openHelperVaHex": "0x004238eb",
        "openHelperCallVaHex": "0x004233cf",
        "createFileCallSites": ["0x00423917", "0x0042393b", "0x0042395f"],
        "readFileCallSites": [row["callVaHex"] for row in read_calls],
        "closeHandleCallVaHex": "0x00423478",
        "selectorGroupLoadVaHex": "0x004234a3",
        "selectorSlotLoadVaHex": "0x004234b1",
        "selectedPointerStoreVaHex": "0x004234ba",
        "selectedPointerGlobalVaHex": "0x0059de30",
        "readBlocks": read_calls,
    }
    slot_patch = {
        "baseString": "savedat0.dat",
        "copyBaseStringVaHex": "0x00423337",
        "slotCharPatchVaHex": "0x0042334f",
        "slotCharOffset": 7,
        "formula": "filename[7] = slot_argument + 0x31",
        "expectedRuntimeNames": ["savedat1.dat", "savedat2.dat", "savedat3.dat", "...", "savedat9.dat"],
    }
    summary = {
        "staticSavePathKnown": True,
        "saveDataDirectoryStringCount": sum(1 for row in strings if row["text"] == "SaveData\\"),
        "slotFilenameStringCount": sum(1 for row in strings if row["text"].startswith("savedat")),
        "registrySavePathStringCount": sum(1 for row in strings if row["text"] == "SavePath"),
        "loadRoutineVaHex": loader["functionVaHex"],
        "openHelperVaHex": loader["openHelperVaHex"],
        "readFileCallCountInLoader": len(read_calls),
        "selectorStoreVaHex": loader["selectedPointerStoreVaHex"],
        "runtimeInputMenuReachabilityProven": False,
        "capturedSelector20SavePresent": False,
        "syntheticSelectorProbeRemainsDiagnostic": True,
        "promotionAllowed": False,
    }
    return {
        "summary": summary,
        "strings": strings,
        "apis": apis,
        "pathBuilders": path_builders,
        "slotFilenamePatch": slot_patch,
        "loader": loader,
        "apiChecks": {
            "readFileHasLoaderCalls": all(
                row["callVaHex"] in {site["vaHex"] for site in api_by_name.get("ReadFile", {}).get("callSites", [])}
                for row in read_calls
            ),
            "createFileHasOpenHelperCalls": all(
                va in {site["vaHex"] for site in api_by_name.get("CreateFileA", {}).get("callSites", [])}
                for va in loader["createFileCallSites"]
            ),
            "findFirstFileHasWildcardCall": "0x004230eb"
            in {site["vaHex"] for site in api_by_name.get("FindFirstFileA", {}).get("callSites", [])},
            "registrySavePathReadPresent": "0x004266fd"
            in {site["vaHex"] for site in api_by_name.get("RegQueryValueExA", {}).get("callSites", [])},
        },
        "verifiedInstructionBytes": verify_expected_bytes(exe, sections),
        "conclusion": {
            "status": "static-save-path-known-runtime-route-proof-missing",
            "detail": (
                "The EXE has a concrete SaveData\\savedatN.dat path builder and the selector loader reads "
                "three blocks before storing 0x0059de30, but no captured gameplay save or runtime input path "
                "currently proves selector 2:0 for map1_01a -> map2_02d."
            ),
            "nextProbe": (
                "A controlled diagnostic can place a backup-restored synthetic or captured savedatN.dat under "
                "SaveData and try to drive the original load menu, but a synthetic file must remain excluded "
                "from route promotion."
            ),
        },
    }


def markdown(summary: dict) -> str:
    s = summary["summary"]
    lines = [
        "# Runtime Save Path Context",
        "",
        "Original EXE save path and save-loader evidence.",
        "",
        "## Conclusion",
        "",
        f"- status: `{summary['conclusion']['status']}`",
        f"- static save path known: `{s['staticSavePathKnown']}`",
        f"- runtime input menu reachability proven: `{s['runtimeInputMenuReachabilityProven']}`",
        f"- promotion allowed: `{s['promotionAllowed']}`",
        f"- detail: {summary['conclusion']['detail']}",
        f"- next probe: {summary['conclusion']['nextProbe']}",
        "",
        "## Save Path Strings",
        "",
        "| string VA | text | direct refs | meaning |",
        "| --- | --- | ---: | --- |",
    ]
    for row in summary["strings"]:
        lines.append(
            f"| `{row['vaHex']}` | `{row['text']}` | {row['directRefCount']} | {row['meaning']} |"
        )

    lines.extend(
        [
            "",
            "## Path Builders",
            "",
            "| name | function | filename | directory | output buffer | evidence |",
            "| --- | --- | --- | --- | --- | --- |",
        ]
    )
    for row in summary["pathBuilders"]:
        lines.append(
            f"| {row['name']} | `{row['functionVaHex']}` | `{row['filenameStringVaHex']}` | "
            f"`{row['directoryStringVaHex']}` | `{row['outputBufferVaHex']}` | {row['apiEvidence']} |"
        )

    patch = summary["slotFilenamePatch"]
    lines.extend(
        [
            "",
            "## Slot Filename Patch",
            "",
            f"- base string: `{patch['baseString']}`",
            f"- copy: `{patch['copyBaseStringVaHex']}`",
            f"- patch: `{patch['slotCharPatchVaHex']}`",
            f"- formula: `{patch['formula']}`",
            f"- expected names: `{', '.join(patch['expectedRuntimeNames'])}`",
            "",
            "## Loader",
            "",
            f"- function: `{summary['loader']['functionVaHex']}`",
            f"- open helper: `{summary['loader']['openHelperVaHex']}` via `{summary['loader']['openHelperCallVaHex']}`",
            f"- selected pointer store: `{summary['loader']['selectedPointerStoreVaHex']}` -> `{summary['loader']['selectedPointerGlobalVaHex']}`",
            "",
            "| ReadFile call | target | save offset | size |",
            "| --- | --- | --- | --- |",
        ]
    )
    for row in summary["loader"]["readBlocks"]:
        lines.append(
            f"| `{row['callVaHex']}` | `{row['targetVaHex']}` | `{row['saveOffsetHex']}` | `{row['sizeHex']}` |"
        )

    lines.extend(
        [
            "",
            "## API Call Sites",
            "",
            "| API | IAT | calls | first call sites |",
            "| --- | --- | ---: | --- |",
        ]
    )
    for row in summary["apis"]:
        sites = ", ".join(f"`{site['vaHex']}`" for site in row["callSites"][:12]) or "-"
        if len(row["callSites"]) > 12:
            sites += f", ... (+{len(row['callSites']) - 12})"
        lines.append(f"| `{row['dll']}!{row['name']}` | `{row['iatVaHex']}` | {row['callCount']} | {sites} |")

    lines.extend(
        [
            "",
            "## API Checks",
            "",
        ]
    )
    for key, value in summary["apiChecks"].items():
        lines.append(f"- {key}: `{value}`")
    return "\n".join(lines).rstrip() + "\n"


def html_page(summary: dict) -> str:
    md = markdown(summary)
    body_lines = []
    in_table = False
    for line in md.splitlines():
        if line.startswith("# "):
            body_lines.append(f"<h1>{html.escape(line[2:])}</h1>")
            continue
        if line.startswith("## "):
            if in_table:
                body_lines.append("</tbody></table>")
                in_table = False
            body_lines.append(f"<h2>{html.escape(line[3:])}</h2>")
            continue
        if line.startswith("| "):
            cells = [cell.strip() for cell in line.strip("|").split("|")]
            if set(cells[0]) == {"-"}:
                continue
            if not in_table:
                body_lines.append("<table><tbody>")
                in_table = True
            tag = "th" if cells and cells[0] in {"string VA", "name", "ReadFile call", "API"} else "td"
            body_lines.append("<tr>" + "".join(f"<{tag}>{html.escape(cell)}</{tag}>" for cell in cells) + "</tr>")
            continue
        if in_table:
            body_lines.append("</tbody></table>")
            in_table = False
        if line.startswith("- "):
            body_lines.append(f"<p>{html.escape(line)}</p>")
        elif line:
            body_lines.append(f"<p>{html.escape(line)}</p>")
    if in_table:
        body_lines.append("</tbody></table>")
    return (
        "<!doctype html><meta charset=\"utf-8\"><title>Runtime Save Path Context</title>"
        "<style>body{font-family:system-ui,sans-serif;margin:24px;line-height:1.45}"
        "table{border-collapse:collapse;margin:12px 0;width:100%;font-size:13px}"
        "td,th{border:1px solid #bbb;padding:4px 6px;text-align:left;vertical-align:top}"
        "code{background:#f3f3f3;padding:1px 3px}</style>"
        + "\n".join(body_lines)
    )


def write_outputs(summary: dict, out_dir: Path = OUT, html_out: Path | None = None) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)
    json_out = out_dir / "runtime_save_path_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.write_text(html_page(summary), encoding="utf-8")
    return json_out


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)
    parser.add_argument(
        "--html-out",
        type=Path,
        default=None,
        help="Optional HTML mirror. Omit to keep the active surface JSON-only.",
    )
    args = parser.parse_args()
    exe = args.exe.read_bytes()
    imports = parse_imports(args.exe)
    summary = build_summary(exe, imports)
    json_out = write_outputs(summary, args.out_dir, args.html_out)
    print(f"wrote runtime save path context -> {json_out}")


if __name__ == "__main__":
    main()