#!/usr/bin/env python3
"""Check whether selector 1:0 rewrites secondaryBranchState after its strongest fill."""
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 read_sections, va_to_offset
from summarize_save_selector_stream_traces import handler_entry


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
PREDECESSOR_ROOT_HEX = "0x00478364"
CURRENT_SELECTOR = "2:0"
CURRENT_READER = 0x00542B0C
FRONTIER_LEAF = 0x00542AE8
SOURCE_RECORD = 0x00542B44
TARGET_RECORD = 0x00542BAC
LEAF_TABLE_START = 0x005429A8
LEAF_TABLE_END = 0x00542A10


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


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


def parse_hex(value: str) -> int:
    return int(value, 16)


def parse_range(value: str) -> tuple[int, int]:
    start, end = value.split("..", 1)
    return parse_hex(start), parse_hex(end)


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 decode_op10(value: int) -> dict | None:
    if (value & 0xFF) != 0x10:
        return None
    stream_plus_1 = (value >> 8) & 0xFF
    helper_arg = (value >> 16) & 0xFF
    stream_plus_3 = (value >> 24) & 0xFF
    state_table = "primaryBranchState" if stream_plus_1 == 0 else "secondaryBranchState"
    helper_valid = helper_arg <= 0x0B
    script_shaped = stream_plus_3 <= 0x02
    return {
        "valueHex": hex32(value),
        "streamPlus1Hex": hex8(stream_plus_1),
        "helperArgumentHex": hex8(helper_arg),
        "streamPlus3Hex": hex8(stream_plus_3),
        "stateTable": state_table,
        "helperDispatchValid": helper_valid,
        "scriptShaped": script_shaped,
        "validSecondaryFill": state_table == "secondaryBranchState" and helper_valid and script_shaped,
    }


def selector_key(row: dict) -> str:
    return f"{row.get('group')}:{row.get('slot')}"


def find_selector_by_root(selectors: list[dict], root_hex: str) -> dict | None:
    return next((row for row in selectors if row.get("selectedPointerHex") == root_hex), None)


def find_selector_by_key(selectors: list[dict], key: str) -> dict | None:
    return next((row for row in selectors if selector_key(row) == key), None)


def next_selector_by_va(selectors: list[dict], root_hex: str) -> dict | None:
    root = parse_hex(root_hex)
    later = [
        row for row in selectors
        if isinstance(row.get("selectedPointer"), int) and row["selectedPointer"] > root
    ]
    later.sort(key=lambda row: row["selectedPointer"])
    return later[0] if later else None


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


def classify_current_route_pointer(value: int, current_root: int | None, current_end: int | None) -> str | None:
    if current_root is not None and value == current_root:
        return "current-root-exact"
    if value == CURRENT_READER:
        return "current-reader"
    if value == FRONTIER_LEAF:
        return "frontier-leaf"
    if value == SOURCE_RECORD:
        return "source-record"
    if value == TARGET_RECORD:
        return "target-record"
    if LEAF_TABLE_START <= value < LEAF_TABLE_END:
        return "leaf-table-window"
    if current_root is None or current_end is None or not current_root <= value < current_end:
        return None
    if value < CURRENT_READER:
        return "current-root-before-reader"
    if CURRENT_READER < value < current_end:
        return "post-reader-current-root"
    return "current-root-other"


def route_bridge_tail_scan(
    exe: bytes,
    sections: list[dict],
    tail_start: int,
    tail_end: int,
    current_root: int | None,
    current_end: int | None,
) -> dict:
    counts = {
        "pointerDwordCount": 0,
        "dataPointerDwordCount": 0,
        "textPointerDwordCount": 0,
        "currentRootRangePointerCount": 0,
        "currentRootExactRefCount": 0,
        "frontierLeafRefCount": 0,
        "currentReaderRefCount": 0,
        "sourceRecordRefCount": 0,
        "targetRecordRefCount": 0,
        "leafTableWindowPointerCount": 0,
        "dataHandlerLowByteCount": 0,
        "branchCapableRowCount": 0,
        "branchToCurrentRouteCount": 0,
    }
    pointer_classes: dict[str, int] = {}
    branch_target_rows = []
    samples = []
    for va in range(tail_start, tail_end, 4):
        value = dword_at(exe, sections, va)
        if value is None:
            continue
        pointer_section = section_name_for_va(sections, value)
        if pointer_section:
            counts["pointerDwordCount"] += 1
            if pointer_section == ".data":
                counts["dataPointerDwordCount"] += 1
            elif pointer_section == ".text":
                counts["textPointerDwordCount"] += 1
        pointer_class = classify_current_route_pointer(value, current_root, current_end)
        if pointer_class:
            counts["currentRootRangePointerCount"] += 1
            pointer_classes[pointer_class] = pointer_classes.get(pointer_class, 0) + 1
            if pointer_class == "current-root-exact":
                counts["currentRootExactRefCount"] += 1
            elif pointer_class == "frontier-leaf":
                counts["frontierLeafRefCount"] += 1
            elif pointer_class == "current-reader":
                counts["currentReaderRefCount"] += 1
            elif pointer_class == "source-record":
                counts["sourceRecordRefCount"] += 1
            elif pointer_class == "target-record":
                counts["targetRecordRefCount"] += 1
            elif pointer_class == "leaf-table-window":
                counts["leafTableWindowPointerCount"] += 1
        opcode = value & 0xFF
        handler = handler_entry(exe, sections, opcode)
        if handler.get("handlerSection") == ".data":
            counts["dataHandlerLowByteCount"] += 1
        if handler.get("canJumpToDwordAtPlus4"):
            counts["branchCapableRowCount"] += 1
            branch_target = dword_at(exe, sections, va + 4)
            target_class = classify_current_route_pointer(branch_target, current_root, current_end) if branch_target is not None else None
            if target_class:
                counts["branchToCurrentRouteCount"] += 1
                if len(branch_target_rows) < 24:
                    branch_target_rows.append({
                        "vaHex": hex32(va),
                        "valueHex": hex32(value),
                        "opcodeHex": hex8(opcode),
                        "handlerVaHex": handler.get("handlerVaHex"),
                        "targetVaHex": hex32(branch_target),
                        "targetClass": target_class,
                    })
        if pointer_class or pointer_section or handler.get("handlerSection") == ".data":
            if len(samples) < 32:
                samples.append({
                    "vaHex": hex32(va),
                    "valueHex": hex32(value),
                    "lowOpcodeHex": hex8(opcode),
                    "pointerSection": pointer_section,
                    "currentRoutePointerClass": pointer_class,
                    "handlerVaHex": handler.get("handlerVaHex"),
                    "handlerSection": handler.get("handlerSection"),
                })
    direct_route_bridge = any(
        counts[key] > 0
        for key in [
            "currentRootExactRefCount",
            "frontierLeafRefCount",
            "currentReaderRefCount",
            "sourceRecordRefCount",
            "targetRecordRefCount",
            "branchToCurrentRouteCount",
        ]
    )
    return {
        "currentSelector": CURRENT_SELECTOR,
        "currentRootHex": hex32(current_root) if current_root is not None else None,
        "currentRootRangeHex": (
            f"{hex32(current_root)}..{hex32(current_end)}"
            if current_root is not None and current_end is not None
            else None
        ),
        "sourceRecordHex": hex32(SOURCE_RECORD),
        "targetRecordHex": hex32(TARGET_RECORD),
        "frontierLeafHex": hex32(FRONTIER_LEAF),
        "currentReaderHex": hex32(CURRENT_READER),
        **counts,
        "currentRoutePointerClasses": dict(sorted(pointer_classes.items())),
        "branchTargetRows": branch_target_rows,
        "sampleRows": samples,
        "directCurrentRouteBridgeFound": direct_route_bridge,
        "classification": "tail-current-route-bridge" if direct_route_bridge else "no-tail-current-route-bridge",
    }


def build_summary(exe: bytes, selectors: list[dict], fill_roots: dict) -> dict:
    sections = read_sections(exe)
    predecessor = next(
        (row for row in fill_roots.get("roots") or [] if row.get("rootHex") == PREDECESSOR_ROOT_HEX),
        None,
    )
    if predecessor is None:
        raise ValueError(f"missing predecessor fill root {PREDECESSOR_ROOT_HEX}")
    root_start, root_end = parse_range(predecessor["rootRangeHex"])
    fills = sorted(predecessor.get("fills") or [], key=lambda row: parse_hex(row["vaHex"]))
    if not fills:
        raise ValueError("predecessor root has no fills")
    last_fill = fills[-1]
    last_fill_va = parse_hex(last_fill["vaHex"])
    tail_start = last_fill_va + 4
    tail_rows = []
    low_opcode_counts: dict[str, int] = {}
    for va in range(tail_start, root_end, 4):
        value = dword_at(exe, sections, va)
        if value is None:
            continue
        opcode_hex = hex8(value & 0xFF)
        low_opcode_counts[opcode_hex] = low_opcode_counts.get(opcode_hex, 0) + 1
        decoded = decode_op10(value)
        row = {
            "vaHex": hex32(va),
            "valueHex": hex32(value),
            "lowOpcodeHex": opcode_hex,
            "op10": decoded,
        }
        if decoded or value in {0, 0x00000240, 0x2A000124} or (0x00400000 <= value <= 0x00600000):
            tail_rows.append(row)
    op10_rows = [row for row in tail_rows if row.get("op10")]
    valid_secondary = [row for row in op10_rows if (row.get("op10") or {}).get("validSecondaryFill")]
    next_selector = next_selector_by_va(selectors, PREDECESSOR_ROOT_HEX)
    predecessor_selector = find_selector_by_root(selectors, PREDECESSOR_ROOT_HEX)
    current_selector = find_selector_by_key(selectors, CURRENT_SELECTOR)
    current_root = parse_hex(current_selector["selectedPointerHex"]) if current_selector else None
    current_next = next_selector_by_va(selectors, current_selector["selectedPointerHex"]) if current_selector else None
    current_end = parse_hex(current_next["selectedPointerHex"]) if current_next else None
    tail_route_bridge = route_bridge_tail_scan(exe, sections, tail_start, root_end, current_root, current_end)
    conclusion = (
        "Within selector root 1:0, the strongest secondaryBranchState fill sites are at 0x004844d0 and "
        "0x004844d8. The tail from 0x004844dc to the next root boundary has no opcode 0x10 rows and no valid "
        "secondaryBranchState helper fill, so this root does not locally reset the predecessor fill after it is written. "
        f"The same tail has {tail_route_bridge['currentRootRangePointerCount']} pointer(s) into the current 2:0 root range, "
        f"but exact current-root/frontier-leaf/current-reader/source-record/target-record refs are "
        f"{tail_route_bridge['currentRootExactRefCount']}/{tail_route_bridge['frontierLeafRefCount']}/"
        f"{tail_route_bridge['currentReaderRefCount']}/{tail_route_bridge['sourceRecordRefCount']}/"
        f"{tail_route_bridge['targetRecordRefCount']}, and branch-capable rows targeting current-route data are "
        f"{tail_route_bridge['branchToCurrentRouteCount']}. "
        "This narrows the persistence gap, but it still does not prove runtime execution order or rule out resets in "
        "untraced VM/helper paths outside this root."
    )
    return {
        "predecessorSelector": selector_key(predecessor_selector) if predecessor_selector else None,
        "predecessorRootHex": PREDECESSOR_ROOT_HEX,
        "predecessorRootRangeHex": predecessor["rootRangeHex"],
        "predecessorFieldMaps": (predecessor_selector or {}).get("fieldMaps") or [],
        "nextRootHex": next_selector.get("selectedPointerHex") if next_selector else hex32(root_end),
        "nextRootSelector": selector_key(next_selector) if next_selector else None,
        "tailRangeHex": f"{hex32(tail_start)}..{hex32(root_end)}",
        "predecessorFillVas": [row["vaHex"] for row in fills],
        "lastFillVaHex": last_fill["vaHex"],
        "lastFillValueHex": last_fill["valueHex"],
        "tailDwordCount": max(0, (root_end - tail_start) // 4),
        "tailOpcode10RowCount": len(op10_rows),
        "tailValidSecondaryFillCount": len(valid_secondary),
        "tailLowOpcodeCounts": dict(sorted(low_opcode_counts.items())),
        "tailRows": tail_rows,
        "tailCurrentRouteBridge": tail_route_bridge,
        "tailCurrentRootRangePointerCount": tail_route_bridge["currentRootRangePointerCount"],
        "tailCurrentRootExactRefCount": tail_route_bridge["currentRootExactRefCount"],
        "tailFrontierLeafRefCount": tail_route_bridge["frontierLeafRefCount"],
        "tailCurrentReaderRefCount": tail_route_bridge["currentReaderRefCount"],
        "tailSourceRecordRefCount": tail_route_bridge["sourceRecordRefCount"],
        "tailTargetRecordRefCount": tail_route_bridge["targetRecordRefCount"],
        "tailBranchToCurrentRouteCount": tail_route_bridge["branchToCurrentRouteCount"],
        "tailDirectCurrentRouteBridgeFound": tail_route_bridge["directCurrentRouteBridgeFound"],
        "localTailResetFound": bool(valid_secondary),
        "promotionStatus": "blocked",
        "remainingProofs": [
            "prove selector 1:0 executes before selector 2:0 in the normal runtime path",
            "prove no global VM/helper reset outside the 1:0 root clears secondaryBranchState before the current reader",
            "find strict map1_01a source coordinate or hotspot",
        ],
        "conclusion": conclusion,
    }


def format_counts(counts: dict[str, int]) -> str:
    return ", ".join(f"{key}:{value}" for key, value in counts.items()) or "-"


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Predecessor Tail Reset",
        "",
        f"- predecessor: `{summary['predecessorSelector']}` root `{summary['predecessorRootHex']}`",
        f"- root range: `{summary['predecessorRootRangeHex']}`",
        f"- next root: `{summary['nextRootSelector']}` `{summary['nextRootHex']}`",
        f"- fill sites: {', '.join(f'`{va}`' for va in summary['predecessorFillVas'])}",
        f"- tail range after last fill: `{summary['tailRangeHex']}`",
        f"- tail opcode 0x10 rows: {summary['tailOpcode10RowCount']}",
        f"- tail valid secondary fills: {summary['tailValidSecondaryFillCount']}",
        f"- local tail reset found: {summary['localTailResetFound']}",
        f"- tail current-root-range pointers: {summary['tailCurrentRootRangePointerCount']}",
        f"- tail direct current-route bridge found: {summary['tailDirectCurrentRouteBridgeFound']}",
        f"- promotion status: {summary['promotionStatus']}",
        "",
        summary["conclusion"],
        "",
        "## Tail Opcode Counts",
        "",
        format_counts(summary["tailLowOpcodeCounts"]),
        "",
        "## Tail Current-Route Bridge Scan",
        "",
        f"- current selector/root: `{summary['tailCurrentRouteBridge']['currentSelector']}` `{summary['tailCurrentRouteBridge']['currentRootRangeHex']}`",
        f"- pointer dwords: {summary['tailCurrentRouteBridge']['pointerDwordCount']} (data {summary['tailCurrentRouteBridge']['dataPointerDwordCount']}, text {summary['tailCurrentRouteBridge']['textPointerDwordCount']})",
        f"- current root range pointers: {summary['tailCurrentRootRangePointerCount']}",
        f"- exact current/frontier/current-reader/source-record/target-record refs: {summary['tailCurrentRootExactRefCount']}/{summary['tailFrontierLeafRefCount']}/{summary['tailCurrentReaderRefCount']}/{summary['tailSourceRecordRefCount']}/{summary['tailTargetRecordRefCount']}",
        f"- leaf-table window pointers: {summary['tailCurrentRouteBridge']['leafTableWindowPointerCount']}",
        f"- branch-capable rows / branch targets into current route: {summary['tailCurrentRouteBridge']['branchCapableRowCount']} / {summary['tailBranchToCurrentRouteCount']}",
        f"- classification: `{summary['tailCurrentRouteBridge']['classification']}`",
        "",
        "| va | value | low opcode | pointer section | current route class | handler |",
        "| --- | --- | --- | --- | --- | --- |",
    ]
    for row in summary["tailCurrentRouteBridge"]["sampleRows"]:
        lines.append(
            f"| `{row['vaHex']}` | `{row['valueHex']}` | `{row['lowOpcodeHex']}` | "
            f"{row.get('pointerSection') or '-'} | {row.get('currentRoutePointerClass') or '-'} | "
            f"`{row.get('handlerVaHex') or '-'}` {row.get('handlerSection') or '-'} |"
        )
    lines.extend([
        "",
        "## Tail Evidence Rows",
        "",
        "| va | value | low opcode | op10 table | op10 helper | valid secondary fill |",
        "| --- | --- | --- | --- | --- | --- |",
    ])
    for row in summary["tailRows"]:
        op10 = row.get("op10") or {}
        lines.append(
            f"| `{row['vaHex']}` | `{row['valueHex']}` | `{row['lowOpcodeHex']}` | "
            f"{op10.get('stateTable') or '-'} | {op10.get('helperArgumentHex') or '-'} | "
            f"{op10.get('validSecondaryFill') if op10 else '-'} |"
        )
    lines.extend(["", "## Remaining Proofs", ""])
    lines.extend(f"- {item}" for item in summary["remainingProofs"])
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    rows = "\n".join(
        "<tr>"
        f"<td><code>{html.escape(row['vaHex'])}</code></td>"
        f"<td><code>{html.escape(row['valueHex'])}</code></td>"
        f"<td><code>{html.escape(row['lowOpcodeHex'])}</code></td>"
        f"<td>{html.escape((row.get('op10') or {}).get('stateTable') or '-')}</td>"
        f"<td>{html.escape((row.get('op10') or {}).get('helperArgumentHex') or '-')}</td>"
        f"<td>{html.escape(str((row.get('op10') or {}).get('validSecondaryFill') if row.get('op10') else '-'))}</td>"
        "</tr>"
        for row in summary["tailRows"]
    )
    proofs = "".join(f"<li>{html.escape(item)}</li>" for item in summary["remainingProofs"])
    bridge = summary["tailCurrentRouteBridge"]
    bridge_rows = "\n".join(
        "<tr>"
        f"<td><code>{html.escape(row['vaHex'])}</code></td>"
        f"<td><code>{html.escape(row['valueHex'])}</code></td>"
        f"<td><code>{html.escape(row['lowOpcodeHex'])}</code></td>"
        f"<td>{html.escape(row.get('pointerSection') or '-')}</td>"
        f"<td>{html.escape(row.get('currentRoutePointerClass') or '-')}</td>"
        f"<td><code>{html.escape(row.get('handlerVaHex') or '-')}</code> {html.escape(row.get('handlerSection') or '-')}</td>"
        "</tr>"
        for row in bridge["sampleRows"]
    )
    return "\n".join([
        "<!doctype html><meta charset=\"utf-8\"><title>Save Selector Predecessor Tail Reset</title>",
        "<style>body{font-family:system-ui,sans-serif;background:#111;color:#eee;max-width:1100px;margin:24px auto}table{border-collapse:collapse}td,th{border:1px solid #444;padding:6px 8px;vertical-align:top}code{color:#9bd4ff}</style>",
        "<h1>Save Selector Predecessor Tail Reset</h1>",
        f"<p>Predecessor <code>{summary['predecessorSelector']}</code> root <code>{summary['predecessorRootHex']}</code>; next root <code>{summary['nextRootSelector']}</code> <code>{summary['nextRootHex']}</code>.</p>",
        f"<p>Tail range <code>{summary['tailRangeHex']}</code>; tail opcode 0x10 rows: {summary['tailOpcode10RowCount']}; tail valid secondary fills: {summary['tailValidSecondaryFillCount']}; local tail reset found: {summary['localTailResetFound']}; tail current-root-range pointers: {summary['tailCurrentRootRangePointerCount']}; tail direct current-route bridge found: {summary['tailDirectCurrentRouteBridgeFound']}; promotion status: {html.escape(summary['promotionStatus'])}</p>",
        f"<p>{html.escape(summary['conclusion'])}</p>",
        f"<p>Tail opcode counts: {html.escape(format_counts(summary['tailLowOpcodeCounts']))}</p>",
        "<h2>Tail Current-Route Bridge Scan</h2>",
        f"<p>Current selector/root: <code>{html.escape(bridge['currentSelector'])}</code> <code>{html.escape(str(bridge['currentRootRangeHex']))}</code>. Pointer dwords: {bridge['pointerDwordCount']} (data {bridge['dataPointerDwordCount']}, text {bridge['textPointerDwordCount']}); current root range pointers: {summary['tailCurrentRootRangePointerCount']}; exact current/frontier/current-reader/source-record/target-record refs: {summary['tailCurrentRootExactRefCount']}/{summary['tailFrontierLeafRefCount']}/{summary['tailCurrentReaderRefCount']}/{summary['tailSourceRecordRefCount']}/{summary['tailTargetRecordRefCount']}; branch targets into current route: {summary['tailBranchToCurrentRouteCount']}; classification: <code>{html.escape(bridge['classification'])}</code>.</p>",
        "<table><thead><tr><th>va</th><th>value</th><th>low opcode</th><th>pointer section</th><th>current route class</th><th>handler</th></tr></thead><tbody>",
        bridge_rows,
        "</tbody></table>",
        "<table><thead><tr><th>va</th><th>value</th><th>low opcode</th><th>op10 table</th><th>op10 helper</th><th>valid secondary fill</th></tr></thead><tbody>",
        rows,
        "</tbody></table>",
        f"<h2>Remaining Proofs</h2><ul>{proofs}</ul>",
    ])


def write_outputs(summary: dict, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_predecessor_tail_reset.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (out_dir / "save_selector_predecessor_tail_reset.html").write_text(html_page(summary), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--exe", type=Path, default=ROOT / "Hwanse2.exe")
    parser.add_argument("--selectors", type=Path, default=OUT / "save_scene_selectors.json")
    parser.add_argument("--fill-roots", type=Path, default=OUT / "save_selector_secondary_fill_roots.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    summary = build_summary(
        args.exe.read_bytes(),
        json.loads(args.selectors.read_text(encoding="utf-8")),
        json.loads(args.fill_roots.read_text(encoding="utf-8")),
    )
    write_outputs(summary, args.out_dir)
    print(f"wrote predecessor tail reset -> {args.out_dir / 'save_selector_predecessor_tail_reset.html'}")


if __name__ == "__main__":
    main()