#!/usr/bin/env python3
"""Group direct primaryBranchState writers used before save-selector branches."""
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_branch_state import build_summary as build_branch_state_summary


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
PRIMARY_BRANCH_STATE = 0x0059E370
CLUSTER_GAP = 0x160
KNOWN_GLOBALS = {
    0x004D2488: "comparisonTable4d2488",
    0x004577A6: "objectCandidateTable4577a6",
    0x0045779A: "objectCandidateTable45779a",
    0x004576EC: "partyCandidateLow4576ec",
    0x004576ED: "partyCandidateHigh4576ed",
    0x004576F8: "partyCandidateTable4576f8",
    0x0059DB30: "runtimeObjectBase59db30",
    0x0059E2A8: "runtimeListState59e2a8",
    0x0059E33E: "runtimeFlag59e33e",
    0x0059E344: "runtimeFlag59e344",
    0x0059E345: "runtimeFlag59e345",
    0x0059E34D: "runtimeFlag59e34d",
}
CLUSTER_NOTES = [
    (
        0x0041DCCC,
        0x0041DF2D,
        "object/stat comparison group",
        "Evaluates indexed runtime objects and writes state 0/1/2 according to object presence and stat/level comparisons.",
    ),
    (
        0x0041E0F2,
        0x0041E1F6,
        "party/object condition group",
        "Uses party/object candidate tables and writes state 0/1/2 from the current runtime party/object condition.",
    ),
    (
        0x0041E390,
        0x0041E431,
        "six-slot object condition group",
        "Loops over a short candidate table and writes state 0/1 for available object slots.",
    ),
    (
        0x0041FB36,
        0x0041FB57,
        "list selection initialization group",
        "Seeds one selected slot with state 1 and clears following slots to 0 for a runtime list selection.",
    ),
]


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


def read_dword_refs(exe: bytes, start_offset: int, end_offset: int) -> list[dict]:
    refs = []
    start = max(0, start_offset)
    end = min(len(exe) - 4, end_offset)
    for off in range(start, end + 1):
        value = struct.unpack_from("<I", exe, off)[0]
        name = KNOWN_GLOBALS.get(value)
        if name:
            refs.append({
                "fileOffsetHex": f"0x{off:06x}",
                "valueHex": hex32(value),
                "name": name,
            })
    seen = set()
    unique = []
    for ref in refs:
        key = (ref["valueHex"], ref["name"])
        if key in seen:
            continue
        seen.add(key)
        unique.append(ref)
    return unique


def cluster_label(start: int, end: int) -> tuple[str, str]:
    for note_start, note_end, label, note in CLUSTER_NOTES:
        if start >= note_start and end <= note_end:
            return label, note
    return "direct branch-state writer group", "Direct writes to primaryBranchState detected near this address range."


def cluster_writers(writers: list[dict]) -> list[list[dict]]:
    ordered = sorted(writers, key=lambda row: int(row["refVaHex"], 16))
    clusters: list[list[dict]] = []
    for row in ordered:
        va = int(row["refVaHex"], 16)
        if not clusters or va - int(clusters[-1][-1]["refVaHex"], 16) > CLUSTER_GAP:
            clusters.append([row])
        else:
            clusters[-1].append(row)
    return clusters


def build_summary(exe: bytes, branch_state: dict | None = None) -> dict:
    sections = read_sections(exe)
    branch_state = branch_state or build_branch_state_summary(exe)
    tables = {table.get("name"): table for table in branch_state.get("tables", [])}
    primary = tables.get("primaryBranchState") or {}
    direct_writers = [
        ref for ref in primary.get("refs", [])
        if ref.get("kind") == "indexedWriteImmediate"
    ]
    clusters = []
    for index, rows in enumerate(cluster_writers(direct_writers), start=1):
        vas = [int(row["refVaHex"], 16) for row in rows]
        start = min(vas)
        end = max(vas)
        start_offset = va_to_offset(sections, start - 0x80)
        end_offset = va_to_offset(sections, end + 0x80)
        source_globals = []
        if start_offset is not None and end_offset is not None:
            source_globals = read_dword_refs(exe, start_offset, end_offset)
        values: dict[str, int] = {}
        for row in rows:
            value = row.get("writeValueHex") or "-"
            values[value] = values.get(value, 0) + 1
        label, note = cluster_label(start, end)
        clusters.append({
            "index": index,
            "label": label,
            "note": note,
            "rangeHex": f"{hex32(start)}..{hex32(end)}",
            "writerCount": len(rows),
            "writeValueCounts": values,
            "sourceGlobals": source_globals,
            "writers": [
                {
                    "vaHex": row["refVaHex"],
                    "writeValueHex": row.get("writeValueHex"),
                    "kind": row.get("kind"),
                }
                for row in rows
            ],
        })
    return {
        "scope": "direct writers of save-selector primaryBranchState inherited by current selector roots",
        "primaryBranchState": {
            "baseVaHex": hex32(PRIMARY_BRANCH_STATE),
            "slots": 12,
            "directWriterCount": len(direct_writers),
            "clusterCount": len(clusters),
        },
        "conclusion": (
            "The current save-selector root does not locally fill the branch-state table before the frontier. "
            "The direct primaryBranchState writers are earlier runtime object/party/list evaluation routines, "
            "so primaryBranchState is inherited from earlier runtime state. "
            "The next playable-progress step is to identify which writer cluster executes before selector root 2:0."
        ),
        "clusters": clusters,
    }


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


def format_globals(globals_: list[dict]) -> str:
    return ", ".join(f"{item['name']} `{item['valueHex']}`" for item in globals_) or "-"


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Branch State Writers",
        "",
        "Direct writers for `primaryBranchState` used by save-selector opcode `0x11` branches.",
        "",
        f"- table: `{summary['primaryBranchState']['baseVaHex']}`",
        f"- direct writer refs: {summary['primaryBranchState']['directWriterCount']}",
        f"- writer clusters: {summary['primaryBranchState']['clusterCount']}",
        "",
        summary["conclusion"],
        "",
        "| cluster | range | writes | source globals | note |",
        "| --- | --- | --- | --- | --- |",
    ]
    for cluster in summary["clusters"]:
        lines.append(
            f"| {cluster['label']} | `{cluster['rangeHex']}` | "
            f"{format_counts(cluster['writeValueCounts'])} | {format_globals(cluster['sourceGlobals'])} | "
            f"{cluster['note']} |"
        )
    lines.extend(["", "## Writer References", ""])
    for cluster in summary["clusters"]:
        lines.extend([
            f"### {cluster['label']} `{cluster['rangeHex']}`",
            "",
            "| writer | value |",
            "| --- | --- |",
        ])
        for writer in cluster["writers"]:
            lines.append(f"| `{writer['vaHex']}` | `{writer.get('writeValueHex') or '-'}` |")
        lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    rows = []
    for cluster in summary["clusters"]:
        rows.append(
            "<tr>"
            f"<td>{html.escape(cluster['label'])}</td>"
            f"<td><code>{html.escape(cluster['rangeHex'])}</code></td>"
            f"<td>{html.escape(format_counts(cluster['writeValueCounts']))}</td>"
            f"<td>{format_globals(cluster['sourceGlobals'])}</td>"
            f"<td>{html.escape(cluster['note'])}</td>"
            "</tr>"
        )
    ref_sections = []
    for cluster in summary["clusters"]:
        writer_rows = []
        for writer in cluster["writers"]:
            writer_rows.append(
                "<tr>"
                f"<td><code>{html.escape(writer['vaHex'])}</code></td>"
                f"<td><code>{html.escape(writer.get('writeValueHex') or '-')}</code></td>"
                "</tr>"
            )
        ref_sections.append(
            f"<h2>{html.escape(cluster['label'])} <code>{html.escape(cluster['rangeHex'])}</code></h2>"
            "<table><thead><tr><th>writer</th><th>value</th></tr></thead>"
            f"<tbody>{''.join(writer_rows)}</tbody></table>"
        )
    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 Branch State Writers</title>",
        "  <style>",
        "    body { margin: 24px; background: #101010; color: #eee; font: 14px system-ui, sans-serif; }",
        "    table { border-collapse: collapse; width: 100%; margin-bottom: 24px; }",
        "    th, td { border: 1px solid #333; padding: 6px 8px; vertical-align: top; }",
        "    th { background: #1d1d1d; }",
        "    code { color: #f5d76e; }",
        "  </style>",
        "</head>",
        "<body>",
        "  <h1>Save Selector Branch State Writers</h1>",
        f"  <p>Table: <code>{html.escape(summary['primaryBranchState']['baseVaHex'])}</code>.</p>",
        f"  <p>{html.escape(summary['conclusion'])}</p>",
        "  <table><thead><tr><th>cluster</th><th>range</th><th>writes</th><th>source globals</th><th>note</th></tr></thead>",
        f"  <tbody>{''.join(rows)}</tbody></table>",
        "\n".join(ref_sections),
        "</body>",
        "</html>",
        "",
    ])


def write_outputs(summary: dict, out_dir: Path = OUT) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_branch_state_writers.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\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("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    summary = build_summary(args.exe.read_bytes())
    write_outputs(summary, args.out_dir)
    print(f"wrote save selector branch state writers -> {args.out_dir / 'save_selector_branch_state_writers.json'}")


if __name__ == "__main__":
    main()