#!/usr/bin/env python3
"""Summarize opcode 0x20 descriptor rows that select runtime object bases."""
from __future__ import annotations

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

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

from probe_exe_scene_tables import find_cns_strings, read_sections, va_to_offset


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
SOURCE = "map1_01a"
TARGET = "map2_02d"
GATE_OFFSETS = {0xE8, 0xEA}
SCAN_DWORDS = 96
SELECTION_OPCODES = {0x10, 0x11, 0x12, 0x13}
BASE_SETTER_OPCODES = {0x28, 0x40, 0x41, 0x42, 0x43, 0x44}
FAILED_OPCODE20_OBJECT_BASE_GATE_IDS = [
    "runtime-context-f2-object-pointer",
    "descriptor-script-gate-row",
    "field-map-or-frontier-reference",
    "specific-gate-base-proof",
]
OPCODE20_OBJECT_BASE_MISSING_EVIDENCE = [
    "runtime context+0xf2 object pointer selected by the active descriptor script",
    "descriptor+4 script row that reads/writes gate offsets 0xe8 or 0xea after the object-base selector",
    "field-map CNS or current-frontier reference after the object-base candidate",
    "specific gate-base proof for map1_01a -> map2_02d",
]
OPCODE20_OBJECT_BASE_EVIDENCE_REFS = [
    {
        "path": "Hwanse2.exe",
        "description": "descriptor+4 script scans and opcode 0x42 handler evidence",
    },
    {
        "path": "out/save_selector_opcode20_descriptor_scripts.json",
        "description": "descriptor script rows and frontier/gate reference scans",
    },
    {
        "path": "out/save_selector_selection_buffer_bases.json",
        "description": "runtime object pointer mode still required for context+0xa8",
    },
]


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


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


def load_json(path: Path, fallback: Any) -> Any:
    if not path.exists():
        return fallback
    return json.loads(path.read_text(encoding="utf-8"))


def parse_hex(value: str | None) -> int | None:
    return int(value, 16) if value else 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 selection_operation(opcode: int) -> str:
    return {
        0x10: "fillBranchStateTable",
        0x11: "readSelectedStateAndBranch",
        0x12: "selectActiveStateSlot",
        0x13: "selectMatchingRuntimeSlot",
    }.get(opcode, "other")


def decode_object_selector_mode(row: dict) -> dict:
    stream_plus_1 = parse_hex(row.get("streamPlus1Hex")) or 0
    stream_plus_2 = parse_hex(row.get("streamPlus2Hex")) or 0
    if stream_plus_1 == 0:
        return {
            "objectSelectorMode": "context+0xf2",
            "effectiveBaseExpression": "dword[0x0059db30 + context[0xf2]*4]",
            "fixedObjectIndex": None,
            "handlerEvidenceVaHex": "0x00406480",
            "runtimeObjectPointerRequired": True,
        }
    if stream_plus_1 == 1:
        return {
            "objectSelectorMode": "stream+2",
            "effectiveBaseExpression": f"dword[0x0059db30 + {stream_plus_2}*4]",
            "fixedObjectIndex": stream_plus_2,
            "handlerEvidenceVaHex": "0x004064a0",
            "runtimeObjectPointerRequired": True,
        }
    return {
        "objectSelectorMode": f"unsupported stream+1 {hex8(stream_plus_1)}",
        "effectiveBaseExpression": "unresolved object pointer table mode",
        "fixedObjectIndex": None,
        "handlerEvidenceVaHex": "-",
        "runtimeObjectPointerRequired": True,
    }


def script_words(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    script_va: int,
    frontier_refs: dict[str, int],
    start_index: int,
    end_index: int,
) -> list[dict]:
    rows = []
    frontier_by_va = {value: label for label, value in frontier_refs.items()}
    for index in range(max(0, start_index), min(SCAN_DWORDS, end_index)):
        row_va = script_va + index * 4
        value = dword_at(exe, sections, row_va)
        if value is None:
            break
        opcode = value & 0xFF
        stream_plus_1 = (value >> 8) & 0xFF
        stream_plus_2 = (value >> 16) & 0xFF
        row: dict[str, Any] = {
            "index": index,
            "vaHex": hex32(row_va),
            "valueHex": hex32(value),
            "opcodeHex": hex8(opcode),
            "streamPlus1Hex": hex8(stream_plus_1),
            "streamPlus2Hex": hex8(stream_plus_2),
            "streamPlus3Hex": hex8((value >> 24) & 0xFF),
            "isSelectionOpcode": opcode in SELECTION_OPCODES,
            "isBaseSetterOpcodeShape": opcode in BASE_SETTER_OPCODES,
            "isPointerDword": va_to_offset(sections, value) is not None,
        }
        if opcode in SELECTION_OPCODES:
            row.update({
                "selectionOperation": selection_operation(opcode),
                "isGateOffset": stream_plus_2 in GATE_OFFSETS,
                "writesSelectionBuffer": opcode in {0x12, 0x13},
                "readsSelectionBuffer": opcode == 0x11,
            })
        if value in strings:
            row["cns"] = strings[value]
        if value in frontier_by_va:
            row["currentFrontierTarget"] = frontier_by_va[value]
        rows.append(row)
    return rows


def scan_after_candidate(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    script_va: int,
    start_index: int,
    frontier_refs: dict[str, int],
) -> dict:
    rows = script_words(exe, sections, strings, script_va, frontier_refs, start_index + 1, SCAN_DWORDS)
    gate_selection_rows = [
        row for row in rows
        if row.get("isSelectionOpcode") and row.get("isGateOffset")
    ]
    field_records = [
        row for row in rows
        if (row.get("cns") or "").startswith("map") and not (row.get("cns") or "").startswith("map_")
    ]
    frontier_hits = [row for row in rows if row.get("currentFrontierTarget")]
    return {
        "gateSelectionRowsAfterCandidate": gate_selection_rows[:12],
        "gateSelectionRowsAfterCandidateCount": len(gate_selection_rows),
        "fieldMapRowsAfterCandidate": field_records[:12],
        "fieldMapRowsAfterCandidateCount": len(field_records),
        "currentFrontierRowsAfterCandidate": frontier_hits[:12],
        "currentFrontierRowsAfterCandidateCount": len(frontier_hits),
    }


def collect_candidates(
    exe: bytes,
    descriptor_scripts: dict,
) -> list[dict]:
    sections = read_sections(exe)
    strings = find_cns_strings(exe, sections)
    frontier_refs = {
        label: parse_hex(value)
        for label, value in (descriptor_scripts.get("frontierRefs") or {}).items()
    }
    frontier_refs = {label: value for label, value in frontier_refs.items() if value is not None}
    candidates = []
    for descriptor in descriptor_scripts.get("descriptorRows") or []:
        script4 = next((script for script in descriptor.get("scripts") or [] if script.get("slot") == "script+4"), {})
        script4_va = parse_hex(script4.get("scriptVaHex"))
        if script4_va is None:
            continue
        context_rows = [
            row for row in script4.get("baseSetterRows") or []
            if row.get("writesContextA8") and not row.get("isPointerDword")
        ]
        for row in context_rows:
            if row.get("opcodeHex") != "0x42":
                continue
            index = row.get("index")
            if not isinstance(index, int):
                continue
            later_context_setters = [
                later for later in script4.get("baseSetterRows") or []
                if later.get("writesContextA8") and isinstance(later.get("index"), int) and later["index"] > index
            ]
            later_non_pointer_setters = [later for later in later_context_setters if not later.get("isPointerDword")]
            after_scan = scan_after_candidate(exe, sections, strings, script4_va, index, frontier_refs)
            candidate = {
                "descriptorIndex": descriptor.get("index"),
                "descriptorVaHex": descriptor.get("descriptorVaHex"),
                "script0LinkedCns": descriptor.get("script0LinkedCns") or [],
                "script4VaHex": script4.get("scriptVaHex"),
                "candidateIndex": index,
                "candidateVaHex": row.get("vaHex"),
                "candidateValueHex": row.get("valueHex"),
                "opcodeHex": row.get("opcodeHex"),
                "streamPlus1Hex": row.get("streamPlus1Hex"),
                "streamPlus2Hex": row.get("streamPlus2Hex"),
                "streamPlus3Hex": row.get("streamPlus3Hex"),
                **decode_object_selector_mode(row),
                "nearbyRows": script_words(exe, sections, strings, script4_va, frontier_refs, index - 3, index + 5),
                "laterContextA8SetterCount": len(later_context_setters),
                "laterPointerLikeContextA8SetterCount": sum(1 for later in later_context_setters if later.get("isPointerDword")),
                "laterNonPointerContextA8SetterCount": len(later_non_pointer_setters),
                "isLastNonPointerContextA8SetterInScript": len(later_non_pointer_setters) == 0,
                "finalContextA8SetterRow": script4.get("lastContextA8SetterRow"),
                "finalNonPointerContextA8SetterRow": script4.get("lastNonPointerContextA8SetterRow"),
                "contributesSpecificGateProof": False,
                **after_scan,
            }
            candidates.append(candidate)
    return candidates


def histogram(values: list[str | None]) -> list[dict]:
    counts: dict[str, int] = {}
    for value in values:
        key = value or "none"
        counts[key] = counts.get(key, 0) + 1
    return [
        {"value": value, "count": count}
        for value, count in sorted(counts.items())
    ]


def build_summary(exe: bytes, descriptor_scripts: dict | None = None, selection_buffer_bases: dict | None = None) -> dict:
    descriptor_scripts = descriptor_scripts if descriptor_scripts is not None else load_json(
        OUT / "save_selector_opcode20_descriptor_scripts.json",
        {},
    )
    selection_buffer_bases = selection_buffer_bases if selection_buffer_bases is not None else load_json(
        OUT / "save_selector_selection_buffer_bases.json",
        {},
    )
    candidates = collect_candidates(exe, descriptor_scripts)
    descriptor_indices = sorted({row["descriptorIndex"] for row in candidates})
    immediate_object_index_candidates = [
        row for row in candidates
        if row.get("objectSelectorMode") == "stream+2"
    ]
    gate_selection_after_count = sum(row["gateSelectionRowsAfterCandidateCount"] for row in candidates)
    field_map_after_count = sum(row["fieldMapRowsAfterCandidateCount"] for row in candidates)
    frontier_after_count = sum(row["currentFrontierRowsAfterCandidateCount"] for row in candidates)
    context_mode_candidates = [row for row in candidates if row.get("objectSelectorMode") == "context+0xf2"]
    conclusion = (
        "The 16 non-pointer descriptor+4 context+0xa8 rows are all opcode 0x42 object-base selectors with "
        "stream+1 == 0, so the handler chooses dword[0x0059db30 + context[0xf2]*4] rather than a fixed stream+2 "
        "object index. Scanning the rest of each descriptor+4 script after these rows finds no 0xe8/0xea "
        "selection-buffer reads or writes, no field-map CNS records, and no current-frontier direct references. "
        "These rows are therefore generic runtime object-base reselections, not a specific gate-base proof for "
        "map1_01a -> map2_02d."
    )
    return {
        "source": SOURCE,
        "target": TARGET,
        "descriptorTableVaHex": descriptor_scripts.get("descriptorTableVaHex"),
        "candidateCount": len(candidates),
        "descriptorCountWithCandidates": len(descriptor_indices),
        "descriptorIndicesWithCandidates": descriptor_indices,
        "opcodeHistogram": histogram([row.get("opcodeHex") for row in candidates]),
        "streamPlus1Histogram": histogram([row.get("streamPlus1Hex") for row in candidates]),
        "objectSelectorModeHistogram": histogram([row.get("objectSelectorMode") for row in candidates]),
        "contextF2ObjectSelectorCount": len(context_mode_candidates),
        "immediateObjectIndexCandidateCount": len(immediate_object_index_candidates),
        "gateSelectionRowsAfterCandidateCount": gate_selection_after_count,
        "fieldMapRowsAfterCandidateCount": field_map_after_count,
        "currentFrontierRowsAfterCandidateCount": frontier_after_count,
        "lastNonPointerCandidateCount": sum(1 for row in candidates if row["isLastNonPointerContextA8SetterInScript"]),
        "anyCandidateContributesSpecificGateProof": any(row["contributesSpecificGateProof"] for row in candidates),
        "selectionBufferBasesRuntimePointerModeStillRequired": selection_buffer_bases.get("runtimePointerModeStillRequired"),
        "handlerEvidence": {
            "opcode42HandlerVaHex": "0x0040644f",
            "contextF2WriteVaHex": "0x00406480",
            "stream2WriteVaHex": "0x004064a0",
            "modeMeaning": "stream+1==0 selects context+0xf2; stream+1==1 selects stream+2.",
        },
        "candidates": candidates,
        "runtimeObjectPointerProofRequired": True,
        "controlPathProofStatus": "blocked",
        "proofFound": False,
        "opcode20ObjectBaseProofFound": False,
        "failedOpcode20ObjectBaseGateIds": FAILED_OPCODE20_OBJECT_BASE_GATE_IDS,
        "missingEvidence": OPCODE20_OBJECT_BASE_MISSING_EVIDENCE,
        "evidenceRefs": OPCODE20_OBJECT_BASE_EVIDENCE_REFS,
        "evidenceRefCount": len(OPCODE20_OBJECT_BASE_EVIDENCE_REFS),
        "promotionStatus": "blocked",
        "conclusion": conclusion,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Opcode 0x20 Object Base Candidates",
        "",
        f"- route: {summary['source']} -> {summary['target']}",
        f"- descriptor table: `{summary['descriptorTableVaHex']}`",
        f"- non-pointer object-base candidates: {summary['candidateCount']}",
        f"- descriptors with candidates: {summary['descriptorCountWithCandidates']}",
        f"- context+0xf2 object selectors: {summary['contextF2ObjectSelectorCount']}",
        f"- fixed stream+2 object selectors: {summary['immediateObjectIndexCandidateCount']}",
        f"- gate `0xe8/0xea` rows after candidates: {summary['gateSelectionRowsAfterCandidateCount']}",
        f"- field-map CNS rows after candidates: {summary['fieldMapRowsAfterCandidateCount']}",
        f"- current-frontier refs after candidates: {summary['currentFrontierRowsAfterCandidateCount']}",
        f"- specific gate proof: {summary['anyCandidateContributesSpecificGateProof']}",
        f"- runtime object pointer proof required: {summary['runtimeObjectPointerProofRequired']}",
        f"- proofFound: `{summary['proofFound']}`",
        f"- opcode20ObjectBaseProofFound: `{summary['opcode20ObjectBaseProofFound']}`",
        f"- promotion status: `{summary['promotionStatus']}`",
        "",
        summary["conclusion"],
        "",
        "## Failed Gates",
        "",
    ]
    lines.extend(f"- `{item}`" for item in summary["failedOpcode20ObjectBaseGateIds"])
    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([
        "",
        "## Handler Evidence",
        "",
        "| handler | context+0xf2 write | stream+2 write | mode |",
        "| --- | --- | --- | --- |",
        (
            f"| `{summary['handlerEvidence']['opcode42HandlerVaHex']}` | "
            f"`{summary['handlerEvidence']['contextF2WriteVaHex']}` | "
            f"`{summary['handlerEvidence']['stream2WriteVaHex']}` | "
            f"{summary['handlerEvidence']['modeMeaning']} |"
        ),
        "",
        "## Candidates",
        "",
        "| descriptor | script+0 CNS | candidate | stream+1 | effective base | gate rows after | field maps after | frontier refs after | last non-pointer | later pointer-like setters |",
        "| ---: | --- | --- | --- | --- | ---: | ---: | ---: | --- | ---: |",
    ])
    for row in summary["candidates"]:
        lines.append(
            f"| {row['descriptorIndex']} | {', '.join(row['script0LinkedCns']) or '-'} | "
            f"`{row['candidateVaHex']}` | `{row['streamPlus1Hex']}` | "
            f"`{row['effectiveBaseExpression']}` | "
            f"{row['gateSelectionRowsAfterCandidateCount']} | "
            f"{row['fieldMapRowsAfterCandidateCount']} | "
            f"{row['currentFrontierRowsAfterCandidateCount']} | "
            f"{row['isLastNonPointerContextA8SetterInScript']} | "
            f"{row['laterPointerLikeContextA8SetterCount']} |"
        )
    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["failedOpcode20ObjectBaseGateIds"]
    )
    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"]
    )
    candidate_rows = []
    for row in summary["candidates"]:
        candidate_rows.append(
            "<tr>"
            f"<td>{row['descriptorIndex']}</td>"
            f"<td>{html.escape(', '.join(row['script0LinkedCns']) or '-')}</td>"
            f"<td><code>{html.escape(row['candidateVaHex'])}</code></td>"
            f"<td><code>{html.escape(row['streamPlus1Hex'])}</code></td>"
            f"<td><code>{html.escape(row['effectiveBaseExpression'])}</code></td>"
            f"<td>{row['gateSelectionRowsAfterCandidateCount']}</td>"
            f"<td>{row['fieldMapRowsAfterCandidateCount']}</td>"
            f"<td>{row['currentFrontierRowsAfterCandidateCount']}</td>"
            f"<td>{row['isLastNonPointerContextA8SetterInScript']}</td>"
            f"<td>{row['laterPointerLikeContextA8SetterCount']}</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 Opcode 0x20 Object Base Candidates</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: 1180px; 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 Opcode 0x20 Object Base Candidates</h1>",
        f"  <p>route {html.escape(summary['source'])} -&gt; {html.escape(summary['target'])}; non-pointer object-base candidates: {summary['candidateCount']}; context+0xf2 object selectors: {summary['contextF2ObjectSelectorCount']}; fixed stream+2 object selectors: {summary['immediateObjectIndexCandidateCount']}; gate rows after candidates {summary['gateSelectionRowsAfterCandidateCount']}; field-map rows after candidates {summary['fieldMapRowsAfterCandidateCount']}; frontier refs after candidates {summary['currentFrontierRowsAfterCandidateCount']}; proofFound <code>{summary['proofFound']}</code>; promotion status <code>{html.escape(summary['promotionStatus'])}</code>.</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>Handler Evidence</h2>",
        "  <table><thead><tr><th>handler</th><th>context+0xf2 write</th><th>stream+2 write</th><th>mode</th></tr></thead><tbody>",
        "  <tr>"
        f"<td><code>{html.escape(summary['handlerEvidence']['opcode42HandlerVaHex'])}</code></td>"
        f"<td><code>{html.escape(summary['handlerEvidence']['contextF2WriteVaHex'])}</code></td>"
        f"<td><code>{html.escape(summary['handlerEvidence']['stream2WriteVaHex'])}</code></td>"
        f"<td>{html.escape(summary['handlerEvidence']['modeMeaning'])}</td>"
        "</tr>",
        "  </tbody></table>",
        "  <h2>Candidates</h2>",
        "  <table><thead><tr><th>descriptor</th><th>script+0 CNS</th><th>candidate</th><th>stream+1</th><th>effective base</th><th>gate rows after</th><th>field maps after</th><th>frontier refs after</th><th>last non-pointer</th><th>later pointer-like setters</th></tr></thead><tbody>",
        *candidate_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_opcode20_object_base_candidates.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (out_dir / "save_selector_opcode20_object_base_candidates.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("--descriptor-scripts", type=Path, default=OUT / "save_selector_opcode20_descriptor_scripts.json")
    parser.add_argument("--selection-buffer-bases", type=Path, default=OUT / "save_selector_selection_buffer_bases.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    summary = build_summary(
        args.exe.read_bytes(),
        load_json(args.descriptor_scripts, {}),
        load_json(args.selection_buffer_bases, {}),
    )
    write_outputs(summary, args.out_dir)
    print(
        "wrote save selector opcode 0x20 object base candidates -> "
        f"{args.out_dir / 'save_selector_opcode20_object_base_candidates.html'}"
    )


if __name__ == "__main__":
    main()
