#!/usr/bin/env python3
"""Summarize secondaryBranchState fill candidates in the current selector root."""
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


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
ROOT_START = 0x00540714
ROOT_END = 0x00543578
FRONTIER_READER = 0x00542B0C
SECONDARY_TABLE = 0x0059E360


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 operand_role(exe: bytes, sections: list[dict], va: int) -> dict | None:
    previous = dword_at(exe, sections, va - 4)
    if previous is None:
        return None
    opcode = previous & 0xFF
    if opcode == 0x11:
        return {
            "kind": "branchTargetOperand",
            "ownerVaHex": hex32(va - 4),
            "meaning": "previous opcode 0x11 reads this dword as branch target operand",
        }
    if opcode in {0x16, 0x22, 0x2F, 0x59, 0x5C}:
        return {
            "kind": "streamOperand",
            "ownerVaHex": hex32(va - 4),
            "meaning": f"previous opcode 0x{opcode:02x} can consume this dword as an operand",
        }
    return None


def decode_op10(value: int) -> dict | None:
    if (value & 0xFF) != 0x10:
        return None
    stream_plus_1 = (value >> 8) & 0xFF
    helper_arg = (value >> 16) & 0xFF
    return {
        "valueHex": hex32(value),
        "streamPlus1Hex": hex8(stream_plus_1),
        "helperArgumentHex": hex8(helper_arg),
        "stateTable": "primaryBranchState" if stream_plus_1 == 0 else "secondaryBranchState",
        "helperDispatchValid": helper_arg <= 0x0B,
    }


def secondary_direct_write_summary(branch_state: dict | None) -> dict:
    if not branch_state:
        return {
            "directWriterCount": None,
            "directWriteValueCounts": {},
            "refKindCounts": {},
            "refs": [],
        }
    tables = {table.get("name"): table for table in branch_state.get("tables") or []}
    secondary = tables.get("secondaryBranchState") or {}
    refs = secondary.get("refs") or []
    direct_writers = [
        ref for ref in refs
        if ref.get("kind") == "indexedWriteImmediate"
    ]
    write_counts: dict[str, int] = {}
    for ref in direct_writers:
        value = ref.get("writeValueHex") or "-"
        write_counts[value] = write_counts.get(value, 0) + 1
    return {
        "directWriterCount": len(direct_writers),
        "directWriteValueCounts": write_counts,
        "refKindCounts": secondary.get("refKindCounts") or {},
        "refs": [
            {
                "refVaHex": ref.get("refVaHex"),
                "kind": ref.get("kind"),
                "writeValueHex": ref.get("writeValueHex"),
            }
            for ref in refs
        ],
    }


def build_summary(exe: bytes, branch_state: dict | None = None) -> dict:
    sections = read_sections(exe)
    candidates = []
    for va in range(ROOT_START, ROOT_END, 4):
        value = dword_at(exe, sections, va)
        if value is None:
            continue
        decoded = decode_op10(value)
        if not decoded:
            continue
        if decoded["stateTable"] != "secondaryBranchState":
            continue
        role = operand_role(exe, sections, va)
        before_frontier = va < FRONTIER_READER
        row = {
            "vaHex": hex32(va),
            **decoded,
            "beforeFrontierReader": before_frontier,
            "operandRole": role,
            "routeRelevance": (
                "candidate before frontier but helper arg is outside 0x00..0x0b"
                if before_frontier and not decoded["helperDispatchValid"]
                else "valid helper fill after frontier; cannot initialize the earlier blocker"
                if not before_frontier and decoded["helperDispatchValid"]
                else "valid helper fill before frontier"
                if before_frontier
                else "secondary opcode10-like dword after frontier"
            ),
        }
        candidates.append(row)
    before = [row for row in candidates if row["beforeFrontierReader"]]
    valid_before = [row for row in before if row["helperDispatchValid"]]
    valid_before_non_operand = [row for row in valid_before if not row.get("operandRole")]
    valid_after = [row for row in candidates if not row["beforeFrontierReader"] and row["helperDispatchValid"]]
    direct_writes = secondary_direct_write_summary(branch_state)
    conclusion = (
        "No valid secondaryBranchState opcode 0x10 fill is present before the 0x00542b0c frontier reader in the current selector root. "
        "The only valid secondary fill-like candidate in this root appears after the frontier, so it cannot initialize the blocker. "
        "The secondaryBranchState table also has no direct indexedWriteImmediate refs in .text; its direct refs are handler reads/base "
        "loads rather than route-specific producers. The map1_01a->map2_02d branch therefore still depends on inherited runtime "
        "secondaryBranchState contents from an earlier helper/root or a still-untraced VM path."
    )
    return {
        "rootRangeHex": f"{hex32(ROOT_START)}..{hex32(ROOT_END)}",
        "frontierReaderHex": hex32(FRONTIER_READER),
        "secondaryBranchStateHex": hex32(SECONDARY_TABLE),
        "helperVaHex": "0x00410c90",
        "helperValidCaseRangeHex": "0x00..0x0b",
        "candidateCount": len(candidates),
        "beforeFrontierCount": len(before),
        "validBeforeFrontierCount": len(valid_before),
        "validBeforeFrontierNonOperandCount": len(valid_before_non_operand),
        "validAfterFrontierCount": len(valid_after),
        "secondaryDirectWrites": direct_writes,
        "candidates": candidates,
        "validAfterFrontier": valid_after,
        "conclusion": conclusion,
    }


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


def markdown(summary: dict) -> str:
    lines = [
        "# Save Selector Secondary Branch-State Sources",
        "",
        f"- root range: `{summary['rootRangeHex']}`",
        f"- frontier reader: `{summary['frontierReaderHex']}`",
        f"- secondaryBranchState: `{summary['secondaryBranchStateHex']}`",
        f"- helper: `{summary['helperVaHex']}` valid cases `{summary['helperValidCaseRangeHex']}`",
        f"- secondary opcode10-like candidates: {summary['candidateCount']}",
        f"- before frontier: {summary['beforeFrontierCount']}",
        f"- valid before frontier: {summary['validBeforeFrontierCount']}",
        f"- valid before frontier and not operand: {summary['validBeforeFrontierNonOperandCount']}",
        f"- valid after frontier: {summary['validAfterFrontierCount']}",
        f"- direct secondaryBranchState writers: {summary['secondaryDirectWrites']['directWriterCount']}",
        "",
        summary["conclusion"],
        "",
        "## Direct Secondary Table Refs",
        "",
        f"- ref kinds: {format_counts(summary['secondaryDirectWrites']['refKindCounts'])}",
        f"- direct write values: {format_counts(summary['secondaryDirectWrites']['directWriteValueCounts'])}",
        "",
        "## Current Root Opcode 0x10 Candidates",
        "",
        "| va | value | stream+1 | helper arg | valid | before frontier | operand role | route relevance |",
        "| --- | --- | --- | --- | --- | --- | --- | --- |",
    ]
    for row in summary["candidates"]:
        role = row.get("operandRole") or {}
        lines.append(
            f"| `{row['vaHex']}` | `{row['valueHex']}` | `{row['streamPlus1Hex']}` | `{row['helperArgumentHex']}` | "
            f"{'yes' if row['helperDispatchValid'] else 'no'} | {'yes' if row['beforeFrontierReader'] else 'no'} | "
            f"{role.get('kind') or '-'} {role.get('ownerVaHex') or ''} | {row['routeRelevance']} |"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    direct = summary["secondaryDirectWrites"]
    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['streamPlus1Hex'])}</code></td>"
        f"<td><code>{html.escape(row['helperArgumentHex'])}</code></td>"
        f"<td>{'yes' if row['helperDispatchValid'] else 'no'}</td>"
        f"<td>{'yes' if row['beforeFrontierReader'] else 'no'}</td>"
        f"<td>{html.escape(str((row.get('operandRole') or {}).get('kind') or '-'))}</td>"
        f"<td>{html.escape(row['routeRelevance'])}</td>"
        "</tr>"
        for row in summary["candidates"]
    )
    return "\n".join([
        "<!doctype html><meta charset=\"utf-8\"><title>Save Selector Secondary Branch-State Sources</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 Branch-State Sources</h1>",
        f"<p>Root range: <code>{html.escape(summary['rootRangeHex'])}</code></p>",
        f"<p>Frontier reader: <code>{html.escape(summary['frontierReaderHex'])}</code>; secondaryBranchState: <code>{html.escape(summary['secondaryBranchStateHex'])}</code>; helper: <code>{html.escape(summary['helperVaHex'])}</code></p>",
        f"<p>Valid before frontier: {summary['validBeforeFrontierCount']}; valid before frontier and not operand: {summary['validBeforeFrontierNonOperandCount']}; valid after frontier: {summary['validAfterFrontierCount']}; direct secondaryBranchState writers: {direct['directWriterCount']}</p>",
        f"<p>Direct secondary table ref kinds: {html.escape(format_counts(direct['refKindCounts']))}; direct write values: {html.escape(format_counts(direct['directWriteValueCounts']))}</p>",
        f"<p>{html.escape(summary['conclusion'])}</p>",
        "<table><thead><tr><th>va</th><th>value</th><th>stream+1</th><th>helper arg</th><th>valid</th><th>before frontier</th><th>operand role</th><th>route relevance</th></tr></thead><tbody>",
        rows,
        "</tbody></table>",
    ])


def write_outputs(summary: dict, out_dir: Path, html_out: Path | None = None) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_secondary_state_sources.json").write_text(
        json.dumps(summary, ensure_ascii=False, separators=(",", ":")) + "\n",
        encoding="utf-8",
    )
    if html_out is not None:
        html_out.parent.mkdir(parents=True, exist_ok=True)
        html_out.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("--branch-state", type=Path, default=OUT / "save_selector_branch_state.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    parser.add_argument("--html-out", type=Path, default=None, help="Optional HTML report output path.")
    args = parser.parse_args()
    branch_state = (
        json.loads(args.branch_state.read_text(encoding="utf-8"))
        if args.branch_state.exists()
        else None
    )
    summary = build_summary(args.exe.read_bytes(), branch_state)
    write_outputs(summary, args.out_dir, args.html_out)
    print(f"wrote secondary branch-state sources -> {args.out_dir / 'save_selector_secondary_state_sources.json'}")


if __name__ == "__main__":
    main()