#!/usr/bin/env python3
"""Scan save-selector script streams for selection-buffer readers/writers."""
from __future__ import annotations

import argparse
import bisect
import html
import json
import struct
from collections import Counter, defaultdict
from pathlib import Path

from probe_exe_scene_tables import find_cns_strings, read_sections, va_to_offset


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


def parse_hex(value: str | None) -> int | None:
    return int(value, 16) if value else None


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 data_section(sections: list[dict]) -> dict:
    for section in sections:
        if section["name"] == ".data":
            return section
    raise ValueError("Hwanse2.exe has no .data section")


def stream_operation(value: int) -> dict | None:
    opcode = value & 0xFF
    if opcode not in {0x11, 0x12, 0x13}:
        return None
    stream_plus_1 = (value >> 8) & 0xFF
    stream_plus_2 = (value >> 16) & 0xFF
    state_table = "primaryBranchState" if stream_plus_1 == 0 else "secondaryBranchState"
    if opcode == 0x11:
        return {
            "operation": "reader",
            "opcodeHex": "0x11",
            "stateTable": state_table,
            "meaning": f"{state_table}[selectionBuffer[{hex8(stream_plus_2)}]] == 1",
        }
    return {
        "operation": "writer",
        "opcodeHex": f"0x{opcode:02x}",
        "stateTable": "active branch-state table" if opcode == 0x12 else "runtime slot match",
        "meaning": (
            f"selectionBuffer[{hex8(stream_plus_2)}] = selected active branch-state slot"
            if opcode == 0x12
            else f"selectionBuffer[{hex8(stream_plus_2)}] = matching runtime slot"
        ),
    }


def build_leaf_index(
    selector_references: list[dict],
    leaf_streams: list[dict],
) -> tuple[list[int], dict[int, dict]]:
    contexts: dict[int, dict] = {}
    for ref in selector_references:
        path = ref.get("pathHex") or []
        if not path:
            continue
        leaf = parse_hex(path[-1])
        if leaf is None:
            continue
        context = contexts.setdefault(
            leaf,
            {
                "leafPointerHex": hex32(leaf),
                "labels": [],
                "resources": [],
                "fieldRecords": [],
                "frontierPairs": [],
            },
        )
        label = ref.get("label")
        if label and label not in context["labels"]:
            context["labels"].append(label)
        resource = ref.get("resource")
        if resource and resource not in context["resources"]:
            context["resources"].append(resource)
        if ref.get("sceneMatched") and ref.get("filename"):
            record = {
                "map": ref.get("resource"),
                "filename": ref.get("filename"),
                "sceneIdHex": ref.get("sceneIdHex"),
            }
            if record not in context["fieldRecords"]:
                context["fieldRecords"].append(record)

    for stream in leaf_streams:
        leaf = parse_hex(stream.get("leafPointerHex"))
        if leaf is None:
            continue
        context = contexts.setdefault(
            leaf,
            {
                "leafPointerHex": hex32(leaf),
                "labels": [],
                "resources": [],
                "fieldRecords": [],
                "frontierPairs": [],
            },
        )
        pair = {
            "source": stream.get("source"),
            "target": stream.get("target"),
            "selector": stream.get("selector"),
        }
        if pair not in context["frontierPairs"]:
            context["frontierPairs"].append(pair)
        for record in stream.get("fieldRecords") or []:
            if record not in context["fieldRecords"]:
                context["fieldRecords"].append(record)
    leaves = sorted(contexts)
    return leaves, contexts


def build_root_index(selector_references: list[dict]) -> tuple[list[int], dict[int, dict]]:
    contexts: dict[int, dict] = {}
    for ref in selector_references:
        root = ref.get("selectedPointer")
        if not isinstance(root, int):
            continue
        context = contexts.setdefault(
            root,
            {
                "selectedPointerHex": hex32(root),
                "labels": [],
                "resources": [],
                "fieldRecords": [],
            },
        )
        label = ref.get("label")
        if label and label not in context["labels"]:
            context["labels"].append(label)
        resource = ref.get("resource")
        if resource and resource not in context["resources"]:
            context["resources"].append(resource)
        if ref.get("kind") == "fieldMap" and ref.get("filename"):
            record = {
                "map": ref.get("resource"),
                "filename": ref.get("filename"),
                "sceneIdHex": ref.get("sceneIdHex"),
            }
            if record not in context["fieldRecords"]:
                context["fieldRecords"].append(record)
    roots = sorted(contexts)
    for index, root in enumerate(roots):
        next_root = roots[index + 1] if index + 1 < len(roots) else None
        contexts[root]["rangeStartHex"] = hex32(root)
        contexts[root]["rangeEndHex"] = hex32(next_root) if next_root is not None else None
    return roots, contexts


def nearest_leaf_context(va: int, leaves: list[int], contexts: dict[int, dict], max_distance: int) -> dict | None:
    index = bisect.bisect_right(leaves, va) - 1
    if index < 0:
        return None
    leaf = leaves[index]
    distance = va - leaf
    if distance < 0 or distance > max_distance:
        return None
    context = dict(contexts[leaf])
    context["relativeOffsetHex"] = f"+0x{distance:03x}"
    context["relativeOffset"] = distance
    return context


def root_context_for_va(va: int, roots: list[int], contexts: dict[int, dict]) -> dict | None:
    index = bisect.bisect_right(roots, va) - 1
    if index < 0:
        return None
    root = roots[index]
    next_root = roots[index + 1] if index + 1 < len(roots) else None
    if next_root is not None and va >= next_root:
        return None
    context = dict(contexts[root])
    context["relativeOffsetHex"] = f"+0x{va - root:03x}"
    context["relativeOffset"] = va - root
    return context


def frontier_branch_lookup(frontier_branches: list[dict]) -> dict[int, list[dict]]:
    lookup: dict[int, list[dict]] = defaultdict(list)
    for edge in frontier_branches:
        for branch in edge.get("branchSteps") or []:
            va = parse_hex(branch.get("streamVaHex"))
            if va is None:
                continue
            lookup[va].append({
                "source": edge.get("source"),
                "target": edge.get("target"),
                "condition": branch.get("condition"),
                "fallthroughVaHex": branch.get("fallthroughVaHex"),
                "branchTargetHex": branch.get("branchTargetHex"),
            })
    return lookup


def scan_operations(
    exe: bytes,
    sections: list[dict],
    selector_references: list[dict],
    leaf_streams: list[dict],
    frontier_branches: list[dict],
    context_window: int,
) -> dict:
    section = data_section(sections)
    strings = find_cns_strings(exe, sections)
    leaves, contexts = build_leaf_index(selector_references, leaf_streams)
    roots, root_contexts = build_root_index(selector_references)
    frontier_lookup = frontier_branch_lookup(frontier_branches)
    counts: Counter[tuple[str, int, str]] = Counter()
    interesting_rows = []

    start = section["raw"]
    end = section["raw"] + section["raw_size"] - 3
    for offset in range(start, end, 4):
        value = struct.unpack_from("<I", exe, offset)[0]
        op = stream_operation(value)
        if not op:
            continue
        va = section["va"] + offset - section["raw"]
        stream_plus_1 = (value >> 8) & 0xFF
        stream_plus_2 = (value >> 16) & 0xFF
        counts[(op["operation"], stream_plus_2, op["stateTable"])] += 1
        if stream_plus_2 != INTERESTING_OFFSET:
            continue

        next_value = dword_at(exe, sections, va + 4)
        row = {
            "vaHex": hex32(va),
            "valueHex": hex32(value),
            "operation": op["operation"],
            "opcodeHex": op["opcodeHex"],
            "streamPlus1Hex": hex8(stream_plus_1),
            "selectionBufferOffsetHex": hex8(stream_plus_2),
            "selectionBufferOffset": stream_plus_2,
            "stateTable": op["stateTable"],
            "meaning": op["meaning"],
            "nextDwordHex": hex32(next_value) if next_value is not None else None,
            "nextDwordCns": strings.get(next_value) if next_value is not None else None,
            "selectorContext": nearest_leaf_context(va, leaves, contexts, context_window),
            "selectorRootContext": root_context_for_va(va, roots, root_contexts),
            "frontierBranches": frontier_lookup.get(va, []),
        }
        if op["operation"] == "reader" and next_value is not None:
            row["branchTargetHex"] = hex32(next_value)
            if next_value in strings:
                row["branchTargetKind"] = f"cns:{strings[next_value]}"
        interesting_rows.append(row)

    writer_rows = [row for row in interesting_rows if row["operation"] == "writer"]
    reader_rows = [row for row in interesting_rows if row["operation"] == "reader"]
    counts_by_offset = [
        {
            "operation": operation,
            "selectionBufferOffsetHex": hex8(selection_offset),
            "selectionBufferOffset": selection_offset,
            "stateTable": state_table,
            "count": count,
        }
        for (operation, selection_offset, state_table), count in sorted(
            counts.items(),
            key=lambda item: (item[0][1], item[0][0], item[0][2]),
        )
    ]
    current_frontier = []
    for row in interesting_rows:
        if row.get("frontierBranches"):
            current_frontier.append(row)
    current_root_hexes = {
        (row.get("selectorRootContext") or {}).get("selectedPointerHex")
        for row in current_frontier
        if row.get("selectorRootContext")
    }
    current_root_hexes.discard(None)
    current_root_rows = [
        row
        for row in interesting_rows
        if (row.get("selectorRootContext") or {}).get("selectedPointerHex") in current_root_hexes
    ]
    first_frontier_va = min(
        (int(row["vaHex"], 16) for row in current_frontier),
        default=None,
    )
    current_root_writers_before_frontier = [
        row
        for row in current_root_rows
        if row.get("operation") == "writer"
        and first_frontier_va is not None
        and int(row["vaHex"], 16) < first_frontier_va
    ]

    return {
        "scan": {
            "section": ".data",
            "alignment": "dword",
            "opcodeHandlers": {
                "0x11": "0x0040b4e6",
                "0x12": "0x0040b55f",
                "0x13": "0x0040b696",
            },
            "selectionBuffer": {
                "contextFieldHex": "0x000000a8",
                "interestingOffsetHex": hex8(INTERESTING_OFFSET),
            },
            "contextWindowBytes": context_window,
        },
        "countsByOffset": counts_by_offset,
        "interestingOffsetRows": interesting_rows,
        "writerCountFor0x20": len(writer_rows),
        "readerCountFor0x20": len(reader_rows),
        "currentFrontierRows": current_frontier,
        "currentFrontierRootHexes": sorted(current_root_hexes),
        "currentFrontierRootRows": current_root_rows,
        "currentFrontierRootWriterCountFor0x20": sum(1 for row in current_root_rows if row["operation"] == "writer"),
        "currentFrontierRootReaderCountFor0x20": sum(1 for row in current_root_rows if row["operation"] == "reader"),
        "currentFrontierRootWritersBeforeFirstReader": current_root_writers_before_frontier,
    }


def row_context(row: dict) -> str:
    context = row.get("selectorContext") or {}
    parts = []
    if context.get("labels"):
        parts.append("labels=" + ",".join(context["labels"][:5]))
    if context.get("resources"):
        parts.append("resources=" + ",".join(context["resources"][:5]))
    fields = context.get("fieldRecords") or []
    if fields:
        parts.append("fields=" + ",".join(str(item.get("map")) for item in fields[:4]))
    if context.get("relativeOffsetHex"):
        parts.append("leaf " + context.get("leafPointerHex", "-") + context["relativeOffsetHex"])
    return "; ".join(parts) or "-"


def row_root_context(row: dict) -> str:
    context = row.get("selectorRootContext") or {}
    if not context:
        return "-"
    fields = context.get("fieldRecords") or []
    parts = [
        "root=" + context.get("selectedPointerHex", "-"),
        "labels=" + ",".join(context.get("labels") or []) if context.get("labels") else "",
        "fields=" + ",".join(str(item.get("map")) for item in fields[:6]) if fields else "",
        context.get("relativeOffsetHex", ""),
    ]
    return "; ".join(part for part in parts if part) or "-"


def markdown(summary: dict) -> str:
    scan = summary["scan"]
    lines = [
        "# Save Selector Selection Writers",
        "",
        "Dword-aligned scan for script opcodes that read/write `context+0xa8` selection-buffer byte offsets.",
        "",
        f"- opcode `0x11`: reader handler `{scan['opcodeHandlers']['0x11']}`.",
        f"- opcode `0x12`: writer handler `{scan['opcodeHandlers']['0x12']}`.",
        f"- opcode `0x13`: runtime-slot writer handler `{scan['opcodeHandlers']['0x13']}`.",
        f"- focus offset: `selectionBuffer[{scan['selectionBuffer']['interestingOffsetHex']}]`.",
        f"- writer candidates for `0x20`: {summary['writerCountFor0x20']}.",
        f"- reader candidates for `0x20`: {summary['readerCountFor0x20']}.",
        f"- current selector-root writers for `0x20`: {summary.get('currentFrontierRootWriterCountFor0x20', 0)}.",
        f"- current selector-root writers before first frontier reader: {len(summary.get('currentFrontierRootWritersBeforeFirstReader') or [])}.",
        "",
        "## Current Frontier",
        "",
        "| va | op | value | meaning | branch | context |",
        "| --- | --- | --- | --- | --- | --- |",
    ]
    for row in summary.get("currentFrontierRows") or []:
        branches = "; ".join(
            f"{item.get('source')}->{item.get('target')} {item.get('condition')}"
            for item in row.get("frontierBranches") or []
        ) or "-"
        lines.append(
            f"| {row['vaHex']} | {row['opcodeHex']} | {row['valueHex']} | "
            f"{row['meaning']} | {branches} | {row_context(row)} |"
        )
    if not summary.get("currentFrontierRows"):
        lines.append("| - | - | - | - | - |")

    lines.extend([
        "",
        "## Current Selector Root 0x20 Operations",
        "",
        "| va | op | value | root context | leaf context |",
        "| --- | --- | --- | --- | --- |",
    ])
    for row in summary.get("currentFrontierRootRows") or []:
        lines.append(
            f"| {row['vaHex']} | {row['opcodeHex']} {row['operation']} | {row['valueHex']} | "
            f"{row_root_context(row)} | {row_context(row)} |"
        )
    if not summary.get("currentFrontierRootRows"):
        lines.append("| - | - | - | - | - |")

    writer_rows = [row for row in summary["interestingOffsetRows"] if row["operation"] == "writer"]
    lines.extend([
        "",
        "## 0x20 Writers",
        "",
        "| va | value | stream+1 | context |",
        "| --- | --- | --- | --- |",
    ])
    for row in writer_rows[:120]:
        lines.append(
            f"| {row['vaHex']} | {row['valueHex']} | {row['streamPlus1Hex']} | {row_context(row)} |"
        )
    if not writer_rows:
        lines.append("| - | - | - | - |")
    elif len(writer_rows) > 120:
        lines.append(f"| ... | ... | ... | {len(writer_rows) - 120} additional rows omitted from markdown |")

    lines.extend([
        "",
        "## Counts By Offset",
        "",
        "| operation | offset | table | count |",
        "| --- | --- | --- | --- |",
    ])
    for item in summary["countsByOffset"]:
        if item["selectionBufferOffset"] > 0x30 and item["selectionBufferOffset"] != INTERESTING_OFFSET:
            continue
        lines.append(
            f"| {item['operation']} | {item['selectionBufferOffsetHex']} | {item['stateTable']} | {item['count']} |"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    def table_rows(rows: list[dict]) -> str:
        rendered = []
        for row in rows:
            branches = "<br>".join(
                html.escape(f"{item.get('source')}->{item.get('target')} {item.get('condition')}")
                for item in row.get("frontierBranches") or []
            ) or "-"
            rendered.append(
                "<tr>"
                f"<td><code>{html.escape(row['vaHex'])}</code></td>"
                f"<td><code>{html.escape(row['opcodeHex'])}</code></td>"
                f"<td><code>{html.escape(row['valueHex'])}</code></td>"
                f"<td>{html.escape(row['operation'])}</td>"
                f"<td>{html.escape(row['meaning'])}</td>"
                f"<td>{branches}</td>"
                f"<td>{html.escape(row_root_context(row))}</td>"
                f"<td>{html.escape(row_context(row))}</td>"
                "</tr>"
            )
        return "\n".join(rendered)

    writer_rows = [row for row in summary["interestingOffsetRows"] if row["operation"] == "writer"]
    reader_rows = [row for row in summary["interestingOffsetRows"] if row["operation"] == "reader"]
    count_rows = "\n".join(
        "<tr>"
        f"<td>{html.escape(item['operation'])}</td>"
        f"<td><code>{html.escape(item['selectionBufferOffsetHex'])}</code></td>"
        f"<td>{html.escape(item['stateTable'])}</td>"
        f"<td>{item['count']}</td>"
        "</tr>"
        for item in summary["countsByOffset"]
        if item["selectionBufferOffset"] <= 0x30 or item["selectionBufferOffset"] == INTERESTING_OFFSET
    )
    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 Selection Writers</title>",
        "  <style>",
        "    body { margin: 24px; background: #101010; color: #eee; font: 14px system-ui, sans-serif; }",
        "    table { border-collapse: collapse; width: 100%; margin: 16px 0 28px; }",
        "    th, td { border: 1px solid #333; padding: 6px 8px; vertical-align: top; }",
        "    th { background: #1d1d1d; position: sticky; top: 0; }",
        "    code { color: #f5d76e; }",
        "  </style>",
        "</head>",
        "<body>",
        "  <h1>Save Selector Selection Writers</h1>",
        "  <p>Dword-aligned scan for script opcodes that read/write <code>context+0xa8</code> selection-buffer byte offsets.</p>",
        "  <p>Handlers: <code>0x11</code> reader <code>0x0040b4e6</code>; <code>0x12</code> writer <code>0x0040b55f</code>; <code>0x13</code> runtime-slot writer <code>0x0040b696</code>.</p>",
        f"  <p>Focus: <code>selectionBuffer[0x20]</code>. Writers: {summary['writerCountFor0x20']}. Readers: {summary['readerCountFor0x20']}.</p>",
        f"  <p>Current selector-root writers: {summary.get('currentFrontierRootWriterCountFor0x20', 0)}. Writers before first frontier reader: {len(summary.get('currentFrontierRootWritersBeforeFirstReader') or [])}.</p>",
        "  <h2>Current Frontier</h2>",
        "  <table><thead><tr><th>va</th><th>op</th><th>value</th><th>kind</th><th>meaning</th><th>branch</th><th>root context</th><th>leaf context</th></tr></thead>",
        f"  <tbody>{table_rows(summary.get('currentFrontierRows') or []) or '<tr><td colspan=\"8\">No frontier rows.</td></tr>'}</tbody></table>",
        "  <h2>Current Selector Root 0x20 Operations</h2>",
        "  <table><thead><tr><th>va</th><th>op</th><th>value</th><th>kind</th><th>meaning</th><th>branch</th><th>root context</th><th>leaf context</th></tr></thead>",
        f"  <tbody>{table_rows(summary.get('currentFrontierRootRows') or []) or '<tr><td colspan=\"8\">No selector-root rows.</td></tr>'}</tbody></table>",
        "  <h2>0x20 Writers</h2>",
        "  <table><thead><tr><th>va</th><th>op</th><th>value</th><th>kind</th><th>meaning</th><th>branch</th><th>root context</th><th>leaf context</th></tr></thead>",
        f"  <tbody>{table_rows(writer_rows[:300]) or '<tr><td colspan=\"8\">No writer rows.</td></tr>'}</tbody></table>",
        "  <h2>0x20 Readers</h2>",
        "  <table><thead><tr><th>va</th><th>op</th><th>value</th><th>kind</th><th>meaning</th><th>branch</th><th>root context</th><th>leaf context</th></tr></thead>",
        f"  <tbody>{table_rows(reader_rows[:300]) or '<tr><td colspan=\"8\">No reader rows.</td></tr>'}</tbody></table>",
        "  <h2>Counts By Offset</h2>",
        "  <table><thead><tr><th>operation</th><th>offset</th><th>table</th><th>count</th></tr></thead>",
        f"  <tbody>{count_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_selection_writers.json").write_text(
        json.dumps(summary, ensure_ascii=False, separators=(",", ":")) + "\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("--selector-references", type=Path, default=OUT / "save_scene_selector_references.json")
    parser.add_argument("--leaf-streams", type=Path, default=OUT / "save_selector_leaf_streams.json")
    parser.add_argument("--frontier-branches", type=Path, default=OUT / "save_selector_frontier_branches.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    parser.add_argument("--context-window", type=int, default=0x300)
    args = parser.parse_args()
    exe = args.exe.read_bytes()
    summary = scan_operations(
        exe,
        read_sections(exe),
        json.loads(args.selector_references.read_text(encoding="utf-8")),
        json.loads(args.leaf_streams.read_text(encoding="utf-8")),
        json.loads(args.frontier_branches.read_text(encoding="utf-8")),
        args.context_window,
    )
    write_outputs(summary, args.out_dir)
    print(
        "wrote selection writer scan "
        f"({summary['writerCountFor0x20']} writers, {summary['readerCountFor0x20']} readers for 0x20) "
        f"-> {args.out_dir / 'save_selector_selection_writers.json'}"
    )


if __name__ == "__main__":
    main()