#!/usr/bin/env python3
"""Find selector roots with opcode-shaped secondaryBranchState fill candidates."""
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 byte_at, handler_entry, u32_at


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
CURRENT_ROOT = 0x00540714
FRONTIER_READER = 0x00542B0C
ROUTE_MAPS = {"map1_01a", "map2_02d", "map2_18d", "map2_09g", "map2_10g", "map2_11g", "map2_12h", "map2_17h", "map2_14j", "map2_15j", "map2_16j"}


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


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


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 offset_to_va(sections: list[dict], offset: int) -> int | None:
    for section in sections:
        raw_start = section["raw"]
        raw_end = raw_start + section["raw_size"]
        if raw_start <= offset < raw_end:
            return section["va"] + (offset - raw_start)
    return None


def root_range(root: int, roots: list[int]) -> tuple[int, int]:
    for candidate in roots:
        if candidate > root:
            return root, candidate
    return root, root + 0x4000


def decode_candidate(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
    if stream_plus_1 == 0:
        return None
    if helper_arg > 0x0B:
        return None
    # This filters obvious pointer/text false positives such as 0x5e59c010 while
    # keeping observed script-shaped values like 0x00000210 and 0x01000210.
    if stream_plus_3 > 0x02:
        return None
    return {
        "valueHex": hex32(value),
        "streamPlus1Hex": hex8(stream_plus_1),
        "helperArgumentHex": hex8(helper_arg),
        "streamPlus3Hex": hex8(stream_plus_3),
    }


def fill_entry_reference_summary(exe: bytes, sections: list[dict], rows: list[dict]) -> dict:
    fill_targets: dict[int, list[int]] = {}
    for index, row in enumerate(rows):
        for fill in row.get("fills") or []:
            fill_va = int(fill["vaHex"], 16)
            fill_targets.setdefault(fill_va, []).append(index)

    refs_by_row: dict[int, list[dict]] = {index: [] for index in range(len(rows))}
    if fill_targets:
        offset = 0
        while offset + 4 <= len(exe):
            value = struct.unpack_from("<I", exe, offset)[0]
            owner_indices = fill_targets.get(value)
            if owner_indices:
                site_va = offset_to_va(sections, offset)
                for owner_index in owner_indices:
                    row = rows[owner_index]
                    root_start, root_end = (int(part, 16) for part in row["rootRangeHex"].split(".."))
                    refs_by_row[owner_index].append({
                        "siteVaHex": hex32(site_va) if site_va is not None else None,
                        "targetVaHex": hex32(value),
                        "inRootRange": site_va is not None and root_start <= site_va < root_end,
                    })
            offset += 1

    branch_targets_by_row: dict[int, list[dict]] = {index: [] for index in range(len(rows))}
    for index, row in enumerate(rows):
        fills = {int(fill["vaHex"], 16) for fill in row.get("fills") or []}
        if not fills:
            continue
        root_start, root_end = (int(part, 16) for part in row["rootRangeHex"].split(".."))
        va = root_start
        while va + 8 <= root_end:
            opcode = byte_at(exe, sections, va)
            handler = handler_entry(exe, sections, opcode) if opcode is not None else {}
            target = u32_at(exe, sections, va + 4) if handler.get("canJumpToDwordAtPlus4") else None
            if target in fills:
                branch_targets_by_row[index].append({
                    "siteVaHex": hex32(va),
                    "opcodeHex": hex8(opcode) if opcode is not None else None,
                    "handlerVaHex": handler.get("handlerVaHex"),
                    "targetVaHex": hex32(target),
                })
            va += 4

    row_summaries = []
    for index, row in enumerate(rows):
        refs = refs_by_row[index]
        root_refs = [ref for ref in refs if ref.get("inRootRange")]
        branch_targets = branch_targets_by_row[index]
        row.update({
            "selector": f"{row.get('group')}:{row.get('slot')}",
            "fillEntryDwordRefCount": len(refs),
            "fillEntryRootRangeDwordRefCount": len(root_refs),
            "fillEntryRootBranchTargetCount": len(branch_targets),
            "fillEntryCandidateFound": bool(refs or branch_targets),
            "fillEntryDwordRefs": refs[:12],
            "fillEntryRootBranchTargets": branch_targets[:12],
        })
        row_summaries.append({
            "selector": row["selector"],
            "rootHex": row.get("rootHex"),
            "rootRangeHex": row.get("rootRangeHex"),
            "fieldMaps": row.get("fieldMaps") or [],
            "routeMapOverlap": row.get("routeMapOverlap") or [],
            "fillCount": row.get("fillCount"),
            "fillEntryDwordRefCount": len(refs),
            "fillEntryRootRangeDwordRefCount": len(root_refs),
            "fillEntryRootBranchTargetCount": len(branch_targets),
            "fillEntryCandidateFound": bool(refs or branch_targets),
            "firstFillEntryDwordRef": refs[0] if refs else None,
            "firstFillEntryRootBranchTarget": branch_targets[0] if branch_targets else None,
        })

    roots_with_entry_refs = [row for row in row_summaries if row["fillEntryCandidateFound"]]
    route_roots_with_entry_refs = [
        row for row in roots_with_entry_refs
        if row.get("routeMapOverlap")
    ]
    return {
        "rootEntryReferenceRows": row_summaries,
        "fillEntryReferenceRootCount": len(roots_with_entry_refs),
        "routeOverlapFillEntryReferenceRootCount": len(route_roots_with_entry_refs),
        "fillEntryReferenceRoots": roots_with_entry_refs,
        "routeOverlapFillEntryReferenceRoots": route_roots_with_entry_refs,
    }


def build_summary(exe: bytes, selectors: list[dict]) -> dict:
    sections = read_sections(exe)
    roots = sorted({row["selectedPointer"] for row in selectors if isinstance(row.get("selectedPointer"), int)})
    rows = []
    for row in selectors:
        root = row.get("selectedPointer")
        if not isinstance(root, int):
            continue
        start, end = root_range(root, roots)
        fills = []
        for va in range(start, end, 4):
            value = dword_at(exe, sections, va)
            if value is None:
                continue
            decoded = decode_candidate(value)
            if not decoded:
                continue
            fills.append({
                "vaHex": hex32(va),
                **decoded,
                "beforeCurrentFrontierReader": root == CURRENT_ROOT and va < FRONTIER_READER,
                "afterCurrentFrontierReader": root == CURRENT_ROOT and va > FRONTIER_READER,
            })
        if not fills:
            continue
        maps = row.get("fieldMaps") or []
        overlap = sorted(set(maps) & ROUTE_MAPS)
        rows.append({
            "group": row.get("group"),
            "slot": row.get("slot"),
            "selector": f"{row.get('group')}:{row.get('slot')}",
            "rootHex": row.get("selectedPointerHex") or hex32(root),
            "rootRangeHex": f"{hex32(start)}..{hex32(end)}",
            "fieldMaps": maps,
            "routeMapOverlap": overlap,
            "fillCount": len(fills),
            "fills": fills,
            "routeRelevance": (
                "current root; no before-frontier fill should be used for promotion"
                if root == CURRENT_ROOT
                else "route-overlap root; possible inherited runtime-state producer candidate"
                if overlap
                else "non-route root"
            ),
        })
    entry_ref_summary = fill_entry_reference_summary(exe, sections, rows)
    current = next((row for row in rows if row["rootHex"] == hex32(CURRENT_ROOT)), None)
    route_overlap_rows = [row for row in rows if row.get("routeMapOverlap")]
    predecessor = next((row for row in route_overlap_rows if row.get("selector") == "1:0"), {})
    entry_ref_roots = entry_ref_summary["fillEntryReferenceRoots"]
    route_entry_ref_roots = entry_ref_summary["routeOverlapFillEntryReferenceRoots"]
    entry_ref_selectors = [
        row.get("selector")
        for row in entry_ref_roots
        if row.get("selector")
    ]
    entry_ref_non_route_only = bool(entry_ref_roots) and not route_entry_ref_roots
    entry_ref_exclusion_status = (
        "non-route-only-entry-reference"
        if entry_ref_non_route_only
        else "route-overlap-entry-reference-present"
        if route_entry_ref_roots
        else "no-entry-reference"
    )
    entry_ref_exclusion_detail = (
        f"entrySelectors={','.join(entry_ref_selectors) or '-'}; "
        f"routeOverlapEntryRoots={len(route_entry_ref_roots)}; "
        f"predecessorRefs={predecessor.get('fillEntryDwordRefCount')}/"
        f"{predecessor.get('fillEntryRootRangeDwordRefCount')}/"
        f"{predecessor.get('fillEntryRootBranchTargetCount')}"
    )
    conclusion = (
        "Opcode-shaped secondaryBranchState fills exist in multiple selector roots, including route-overlap roots. "
        "For the current root 0x00540714, the only opcode-shaped secondary fill is after the 0x00542b0c frontier reader, "
        "so it cannot initialize the map1_01a->map2_02d branch. A global fill-entry scan finds no direct dword refs "
        "or root-range branch targets into any route-overlap fill fragment, including the 1:0 predecessor; only non-route "
        "fill roots have an entry-like direct reference. The only entry-ref selector is 55:0, which has no route-map "
        "overlap. The next proof step is therefore runtime/control-flow evidence, "
        "not a static route-overlap fill-site pointer."
    )
    return {
        "currentRootHex": hex32(CURRENT_ROOT),
        "frontierReaderHex": hex32(FRONTIER_READER),
        "routeMaps": sorted(ROUTE_MAPS),
        "rootCount": len(rows),
        "routeOverlapRootCount": len(route_overlap_rows),
        "fillEntryReferenceRootCount": entry_ref_summary["fillEntryReferenceRootCount"],
        "routeOverlapFillEntryReferenceRootCount": entry_ref_summary[
            "routeOverlapFillEntryReferenceRootCount"
        ],
        "fillEntryReferenceRoots": entry_ref_summary["fillEntryReferenceRoots"],
        "routeOverlapFillEntryReferenceRoots": entry_ref_summary["routeOverlapFillEntryReferenceRoots"],
        "fillEntryReferenceSelectors": entry_ref_selectors,
        "fillEntryReferenceNonRouteOnly": entry_ref_non_route_only,
        "fillEntryReferenceExclusionStatus": entry_ref_exclusion_status,
        "fillEntryReferenceExclusionDetail": entry_ref_exclusion_detail,
        "predecessorFillEntryCandidateFound": predecessor.get("fillEntryCandidateFound"),
        "predecessorFillEntryDwordRefCount": predecessor.get("fillEntryDwordRefCount"),
        "predecessorFillEntryRootRangeDwordRefCount": predecessor.get(
            "fillEntryRootRangeDwordRefCount"
        ),
        "predecessorFillEntryRootBranchTargetCount": predecessor.get(
            "fillEntryRootBranchTargetCount"
        ),
        "currentRoot": current,
        "roots": rows,
        "routeOverlapRoots": route_overlap_rows,
        "conclusion": conclusion,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Secondary Fill Roots",
        "",
        f"- current root: `{summary['currentRootHex']}`",
        f"- frontier reader: `{summary['frontierReaderHex']}`",
        f"- roots with opcode-shaped secondary fills: {summary['rootCount']}",
        f"- route-overlap roots: {summary['routeOverlapRootCount']}",
        f"- roots with fill-entry refs/branch targets: {summary['fillEntryReferenceRootCount']}",
        f"- route-overlap roots with fill-entry refs/branch targets: {summary['routeOverlapFillEntryReferenceRootCount']}",
        f"- fill-entry ref selectors: `{','.join(summary.get('fillEntryReferenceSelectors') or []) or '-'}`",
        f"- fill-entry refs restricted to non-route roots: {summary.get('fillEntryReferenceNonRouteOnly')}",
        f"- fill-entry ref exclusion: `{summary.get('fillEntryReferenceExclusionStatus')}` ({summary.get('fillEntryReferenceExclusionDetail')})",
        f"- predecessor fill-entry candidate found: {summary['predecessorFillEntryCandidateFound']}",
        "",
        summary["conclusion"],
        "",
        "| group | slot | root | fills | route overlap | entry refs | root refs | branch targets | relevance | first fills |",
        "| ---: | ---: | --- | ---: | --- | ---: | ---: | ---: | --- | --- |",
    ]
    for row in summary["roots"]:
        fills = ", ".join(f"`{item['vaHex']}={item['valueHex']}`" for item in row["fills"][:5])
        overlap = ", ".join(row.get("routeMapOverlap") or []) or "-"
        lines.append(
            f"| {row.get('group')} | {row.get('slot')} | `{row['rootHex']}` | {row['fillCount']} | "
            f"{overlap} | {row['fillEntryDwordRefCount']} | "
            f"{row['fillEntryRootRangeDwordRefCount']} | {row['fillEntryRootBranchTargetCount']} | "
            f"{row['routeRelevance']} | {fills} |"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    rows = "\n".join(
        "<tr>"
        f"<td>{row.get('group')}</td>"
        f"<td>{row.get('slot')}</td>"
        f"<td><code>{html.escape(row['rootHex'])}</code></td>"
        f"<td>{row['fillCount']}</td>"
        f"<td>{html.escape(', '.join(row.get('routeMapOverlap') or []) or '-')}</td>"
        f"<td>{row['fillEntryDwordRefCount']}</td>"
        f"<td>{row['fillEntryRootRangeDwordRefCount']}</td>"
        f"<td>{row['fillEntryRootBranchTargetCount']}</td>"
        f"<td>{html.escape(row['routeRelevance'])}</td>"
        f"<td>{', '.join('<code>' + html.escape(item['vaHex']) + '=' + html.escape(item['valueHex']) + '</code>' for item in row['fills'][:5])}</td>"
        "</tr>"
        for row in summary["roots"]
    )
    return "\n".join([
        "<!doctype html><meta charset=\"utf-8\"><title>Save Selector Secondary Fill Roots</title>",
        "<style>body{font-family:system-ui,sans-serif;background:#111;color:#eee}table{border-collapse:collapse}td,th{border:1px solid #444;padding:6px 8px;vertical-align:top}code{color:#9bd4ff}</style>",
        "<h1>Save Selector Secondary Fill Roots</h1>",
        f"<p>Current root: <code>{summary['currentRootHex']}</code>; frontier reader: <code>{summary['frontierReaderHex']}</code></p>",
        f"<p>Fill-entry refs/branch targets: {summary['fillEntryReferenceRootCount']} roots; route-overlap roots: {summary['routeOverlapFillEntryReferenceRootCount']}; predecessor candidate: {summary['predecessorFillEntryCandidateFound']}.</p>",
        f"<p>Fill-entry ref exclusion: <code>{html.escape(str(summary.get('fillEntryReferenceExclusionStatus')))}</code>; selectors <code>{html.escape(','.join(summary.get('fillEntryReferenceSelectors') or []) or '-')}</code>; non-route only {summary.get('fillEntryReferenceNonRouteOnly')}.</p>",
        f"<p>{html.escape(summary['conclusion'])}</p>",
        "<table><thead><tr><th>group</th><th>slot</th><th>root</th><th>fills</th><th>route overlap</th><th>entry refs</th><th>root refs</th><th>branch targets</th><th>relevance</th><th>first fills</th></tr></thead><tbody>",
        rows,
        "</tbody></table>",
    ])


def write_outputs(summary: dict, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_secondary_fill_roots.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (out_dir / "save_selector_secondary_fill_roots.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("--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")),
    )
    write_outputs(summary, args.out_dir)
    print(f"wrote secondary fill roots -> {args.out_dir / 'save_selector_secondary_fill_roots.html'}")


if __name__ == "__main__":
    main()