#!/usr/bin/env python3
"""Trace the current save-selector root writers that feed selectionBuffer[0x20]."""
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 find_cns_strings, read_sections, va_to_offset
from summarize_save_selector_stream_traces import trace_stream


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

EVIDENCE_REFS = [
    {
        "path": "out/save_selector_selection_writers.json",
        "fields": [
            "currentFrontierRootWritersBeforeFirstReader",
            "selectorRootContext",
            "selectionBufferOffsetHex",
        ],
    },
    {
        "path": "Hwanse2.exe",
        "fields": [
            "writer stream dwords",
            "opcode handler table",
            "activation context dwords",
        ],
    },
]
FAILED_CURRENT_WRITER_PATH_GATE_IDS = [
    "source-predecessor-current-producer",
    "non-current-selected-pointer-store",
    "runtime-branch-state-slot-proof",
    "selected-root-execution",
]
CURRENT_WRITER_PATH_MISSING_EVIDENCE = [
    "source or predecessor opcode path that produces current root 2:0 before route entry",
    "non-current opcode 0x09/store path writing current root/range into selected pointer 0x0059de30",
    "runtime branch-state slot values proving current local writers choose the route slot on the normal route path",
    "selected-root execution proof tying current writer paths to an external trigger",
]


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


def stream_plus(value_hex: str, index: int) -> int:
    return (int(value_hex, 16) >> (index * 8)) & 0xFF


def opcode_meaning(step: dict) -> str:
    opcode = step.get("opcodeHex")
    value_hex = step.get("valueHex") or "0x00000000"
    stream1 = stream_plus(value_hex, 1)
    stream2 = stream_plus(value_hex, 2)
    table = "primaryBranchState" if stream1 == 0 else "secondaryBranchState"
    if opcode == "0x10":
        if stream2 > 0x0B:
            return (
                f"call helper 0x00410c90 for {table} with out-of-range argument 0x{stream2:02x}; "
                "helper dispatch only handles 0x00..0x0b, so the table is left as inherited runtime state"
            )
        return f"fill {table} using helper argument 0x{stream2:02x}"
    if opcode == "0x11":
        return f"read {table}[selectionBuffer[0x{stream2:02x}]] and gate next stream value"
    if opcode == "0x12":
        return f"write selected active branch-state slot to selectionBuffer[0x{stream2:02x}]"
    if opcode == "0x13":
        return f"write matching party/runtime slot to selectionBuffer[0x{stream2:02x}]"
    if opcode == "0x09":
        mode = stream_plus(value_hex, 1)
        return f"store runtime stream pointer mode 0x{mode:02x} into 0x0059de30"
    if opcode == "0x20":
        return "dispatch/load runtime object or party-member script references"
    return ""


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 context_word(exe: bytes, sections: list[dict], strings: dict[int, str], va: int) -> dict | None:
    value = dword_at(exe, sections, va)
    if value is None:
        return None
    opcode = value & 0xFF
    row = {
        "vaHex": f"0x{va:08x}",
        "valueHex": f"0x{value:08x}",
        "opcodeHex": f"0x{opcode:02x}",
    }
    if value in strings:
        row["cns"] = strings[value]
    elif va_to_offset(sections, value) is not None:
        row["pointerHex"] = f"0x{value:08x}"
    row["meaning"] = opcode_meaning(row)
    return row


def activation_context(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    stream_start: int,
    writer_va: int,
) -> dict:
    rows = []
    for va in range(max(0, stream_start - 0x40), writer_va, 4):
        row = context_word(exe, sections, strings, va)
        if not row:
            continue
        if row.get("meaning") or row.get("pointerHex") or row.get("cns"):
            rows.append(row)
    activators = []
    for row in rows:
        va = parse_hex(row.get("vaHex"))
        value = parse_hex(row.get("valueHex"))
        if va is None or value is None:
            continue
        opcode = value & 0xFF
        mode = (value >> 8) & 0xFF
        if opcode == 0x09 and mode == 0 and va + 4 == stream_start:
            activators.append({
                "vaHex": row["vaHex"],
                "valueHex": row["valueHex"],
                "kind": "opcode09Mode0NextStream",
                "meaning": "opcode 0x09 mode 0 advances to the next dword and stores that stream pointer in 0x0059de30",
            })
        if row.get("pointerHex") == f"0x{stream_start:08x}":
            activators.append({
                "vaHex": row["vaHex"],
                "valueHex": row["valueHex"],
                "kind": "nearbyPointerToStreamStart",
                "meaning": "nearby data word points at this stream start",
            })
    return {
        "contextRows": rows,
        "activators": activators,
    }


def choose_trace_start(exe: bytes, sections: list[dict], strings: dict[int, str], writer_va: int) -> tuple[int, list[dict]]:
    best_start = writer_va
    best_trace = trace_stream(exe, sections, strings, writer_va, 20)
    for start in range(writer_va - 0x30, writer_va + 1, 4):
        if start < 0:
            continue
        trace = trace_stream(exe, sections, strings, start, 20)
        if any(item.get("vaHex") == f"0x{writer_va:08x}" for item in trace):
            best_start = start
            best_trace = trace
            break
    return best_start, best_trace


def field_maps(context: dict | None) -> list[str]:
    maps = []
    for record in (context or {}).get("fieldRecords") or []:
        name = record.get("map")
        if name and name not in maps:
            maps.append(name)
    return maps


def build_rows(exe: bytes, selection_writers: dict) -> list[dict]:
    sections = read_sections(exe)
    strings = find_cns_strings(exe, sections)
    writer_rows = selection_writers.get("currentFrontierRootWritersBeforeFirstReader") or []
    rows = []
    for writer in writer_rows:
        writer_va = parse_hex(writer.get("vaHex"))
        if writer_va is None:
            continue
        writer_value = parse_hex(writer.get("valueHex")) or 0
        writer_opcode = writer_value & 0xFF
        writer_algorithm = (
            "0x13 clears selectionBuffer[stream+2], scans 12 runtime slots, compares 0x0059e344 bytes "
            "with slot fields +0x4a/+0x48, and stores the matching slot index back to selectionBuffer[stream+2]."
            if writer_opcode == 0x13
            else (
                "0x12 selects primaryBranchState when stream+1 is 0, otherwise secondaryBranchState. "
                "If runtime flag 0x00457744 is set it starts from the existing selectionBuffer[stream+2], "
                "then searches the 12-slot table for a nonzero state and stores that slot index back to "
                "selectionBuffer[stream+2]."
            )
        )
        trace_start, trace = choose_trace_start(exe, sections, strings, writer_va)
        trace_rows = []
        for step in trace:
            item = dict(step)
            item["meaning"] = opcode_meaning(step)
            trace_rows.append(item)
        writer_index = next(
            (index for index, step in enumerate(trace_rows) if step.get("vaHex") == writer.get("vaHex")),
            None,
        )
        setup_steps = trace_rows[:writer_index] if writer_index is not None else []
        following_steps = trace_rows[writer_index + 1 :] if writer_index is not None else []
        activation = activation_context(exe, sections, strings, trace_start, writer_va)
        rows.append({
            "writerVaHex": writer.get("vaHex"),
            "writerValueHex": writer.get("valueHex"),
            "streamStartHex": f"0x{trace_start:08x}",
            "selectionBufferOffsetHex": writer.get("selectionBufferOffsetHex"),
            "stateTable": writer.get("stateTable"),
            "rootHex": (writer.get("selectorRootContext") or {}).get("selectedPointerHex"),
            "rootLabels": (writer.get("selectorRootContext") or {}).get("labels") or [],
            "rootFieldMaps": field_maps(writer.get("selectorRootContext")),
            "evidenceRefs": EVIDENCE_REFS,
            "evidenceRefCount": len(EVIDENCE_REFS),
            "proofFound": False,
            "currentWriterPathProofFound": False,
            "failedCurrentWriterPathGateIds": FAILED_CURRENT_WRITER_PATH_GATE_IDS,
            "missingEvidence": CURRENT_WRITER_PATH_MISSING_EVIDENCE,
            "remainingProofs": CURRENT_WRITER_PATH_MISSING_EVIDENCE,
            "promotionStatus": "blocked-current-internal-only",
            "setupSteps": setup_steps,
            "followingSteps": following_steps,
            "trace": trace_rows,
            "activationContext": activation,
            "writerAlgorithm": writer_algorithm,
            "opcode12Algorithm": (
                "0x12 selects primaryBranchState when stream+1 is 0, otherwise secondaryBranchState. "
                "If runtime flag 0x00457744 is set it starts from the existing selectionBuffer[stream+2], "
                "then searches the 12-slot table for a nonzero state and stores that slot index back to "
                "selectionBuffer[stream+2]."
            ),
            "opcode13Algorithm": (
                "0x13 clears selectionBuffer[stream+2], scans 12 runtime slots, compares 0x0059e344 bytes "
                "with slot fields +0x4a/+0x48, and stores the matching slot index back to selectionBuffer[stream+2]."
            ),
            "opcode10HelperRule": (
                "0x10 passes stream+2 to helper 0x00410c90. The helper dispatches only argument values "
                "0x00..0x0b; out-of-range values such as 0x20 return without filling the 12-slot table."
            ),
            "classification": (
                "primary local writer before current frontier reader"
                if writer.get("valueHex") == "0x00200012" and writer_opcode == 0x12
                else "runtime-slot local writer before current frontier reader"
                if writer_opcode == 0x13
                else "secondary local writer before current frontier reader"
            ),
            "nextQuestion": (
                "Treat this 0x13 value as an intermediate overwrite, then continue to the later 0x12 writer."
                if writer_opcode == 0x13
                else (
                    "Resolve the branch-state table values produced by the preceding 0x10 helper calls, "
                    "then identify which slot 0x12 stores into selectionBuffer[0x20]."
                )
            ),
        })
    return rows


def markdown(rows: list[dict]) -> str:
    lines = [
        "# Save Selector Current Writer Paths",
        "",
        "Local traces for the current selector-root writers that appear before the `map1_01a -> map2_02d` frontier reader.",
        "",
        f"- evidence refs: `{len(EVIDENCE_REFS)}`",
        "- promotion status: `blocked-current-internal-only`",
        "- proof found: False",
        f"- failed current-writer gates: {', '.join(FAILED_CURRENT_WRITER_PATH_GATE_IDS)}",
        f"- missing evidence count: {len(CURRENT_WRITER_PATH_MISSING_EVIDENCE)}",
        "",
        "## Evidence Refs",
        "",
    ]
    for ref in EVIDENCE_REFS:
        lines.append(f"- `{ref.get('path')}`: {', '.join(ref.get('fields') or [])}")
    lines.extend(["", "## Missing Evidence", ""])
    lines.extend(f"- {item}" for item in CURRENT_WRITER_PATH_MISSING_EVIDENCE)
    lines.extend([
        "",
        "## Writer Summary",
        "",
        "| writer | stream start | activation | root | maps | setup | following | next question |",
        "| --- | --- | --- | --- | --- | --- | --- | --- |",
    ])
    for row in rows:
        setup = ", ".join(
            f"{step['vaHex']} {step.get('opcodeHex')} {step.get('meaning')}"
            for step in row.get("setupSteps") or []
            if step.get("meaning")
        ) or "-"
        following = ", ".join(
            f"{step['vaHex']} {step.get('opcodeHex')} {step.get('meaning') or step.get('stopReason') or ''}".strip()
            for step in row.get("followingSteps") or []
            if step.get("meaning") or step.get("stopReason")
        ) or "-"
        activation = ", ".join(
            f"{item['vaHex']} {item['kind']}"
            for item in (row.get("activationContext") or {}).get("activators") or []
        ) or "-"
        lines.append(
            f"| {row['writerVaHex']} `{row['writerValueHex']}` | {row['streamStartHex']} | "
            f"{activation} | "
            f"{row.get('rootHex')} {','.join(row.get('rootLabels') or [])} | "
            f"{', '.join(row.get('rootFieldMaps') or []) or '-'} | {setup} | {following} | "
            f"{row['nextQuestion']} |"
        )
    if not rows:
        lines.append("| - | - | - | - | - | - | - | - |")
    lines.extend(["", "## Full Traces", ""])
    for row in rows:
        lines.extend([
            f"### {row['writerVaHex']}",
            "",
            row.get("writerAlgorithm") or row.get("opcode12Algorithm") or "",
            "",
            row.get("opcode10HelperRule") or "",
            "",
            "Activation context:",
            "",
            "| va | value | opcode | meaning | pointer/cns |",
            "| --- | --- | --- | --- | --- |",
        ])
        for item in (row.get("activationContext") or {}).get("contextRows") or []:
            pointer = item.get("pointerHex") or item.get("cns") or "-"
            lines.append(
                f"| {item.get('vaHex')} | `{item.get('valueHex')}` | `{item.get('opcodeHex')}` | "
                f"{item.get('meaning') or '-'} | {pointer} |"
            )
        if not (row.get("activationContext") or {}).get("contextRows"):
            lines.append("| - | - | - | - | - |")
        lines.extend([
            "",
            "| step | va | value | opcode | handler | meaning | stop |",
            "| ---: | --- | --- | --- | --- | --- | --- |",
        ])
        for step in row.get("trace") or []:
            lines.append(
                f"| {step.get('step')} | {step.get('vaHex')} | `{step.get('valueHex')}` | "
                f"`{step.get('opcodeHex')}` | `{step.get('handlerVaHex')}` | "
                f"{step.get('meaning') or '-'} | {step.get('stopReason') or '-'} |"
            )
        lines.append("")
    return "\n".join(lines)


def html_page(rows: list[dict]) -> str:
    evidence_refs = "".join(
        "<li>"
        f"<code>{html.escape(str(ref.get('path')))}</code>: "
        f"{html.escape(', '.join(ref.get('fields') or []))}"
        "</li>"
        for ref in EVIDENCE_REFS
    )
    summary_rows = []
    trace_sections = []
    for row in rows:
        setup = "<br>".join(
            html.escape(f"{step['vaHex']} {step.get('opcodeHex')} {step.get('meaning')}")
            for step in row.get("setupSteps") or []
            if step.get("meaning")
        ) or "-"
        following = "<br>".join(
            html.escape(f"{step['vaHex']} {step.get('opcodeHex')} {step.get('meaning') or step.get('stopReason') or ''}".strip())
            for step in row.get("followingSteps") or []
            if step.get("meaning") or step.get("stopReason")
        ) or "-"
        activation = "<br>".join(
            html.escape(f"{item['vaHex']} {item['kind']}: {item['meaning']}")
            for item in (row.get("activationContext") or {}).get("activators") or []
        ) or "-"
        summary_rows.append(
            "<tr>"
            f"<td><code>{html.escape(row['writerVaHex'])}</code><br><code>{html.escape(row['writerValueHex'])}</code></td>"
            f"<td><code>{html.escape(row['streamStartHex'])}</code></td>"
            f"<td>{activation}</td>"
            f"<td><code>{html.escape(row.get('rootHex') or '-')}</code><br>{html.escape(','.join(row.get('rootLabels') or []))}</td>"
            f"<td>{html.escape(', '.join(row.get('rootFieldMaps') or []) or '-')}</td>"
            f"<td>{setup}</td>"
            f"<td>{following}</td>"
            f"<td>{html.escape(row['nextQuestion'])}</td>"
            "</tr>"
        )
        trace_body = []
        activation_rows = []
        for item in (row.get("activationContext") or {}).get("contextRows") or []:
            pointer = item.get("pointerHex") or item.get("cns") or "-"
            activation_rows.append(
                "<tr>"
                f"<td><code>{html.escape(item.get('vaHex') or '-')}</code></td>"
                f"<td><code>{html.escape(item.get('valueHex') or '-')}</code></td>"
                f"<td><code>{html.escape(item.get('opcodeHex') or '-')}</code></td>"
                f"<td>{html.escape(item.get('meaning') or '-')}</td>"
                f"<td>{html.escape(pointer)}</td>"
                "</tr>"
            )
        for step in row.get("trace") or []:
            trace_body.append(
                "<tr>"
                f"<td>{step.get('step')}</td>"
                f"<td><code>{html.escape(step.get('vaHex') or '-')}</code></td>"
                f"<td><code>{html.escape(step.get('valueHex') or '-')}</code></td>"
                f"<td><code>{html.escape(step.get('opcodeHex') or '-')}</code></td>"
                f"<td><code>{html.escape(step.get('handlerVaHex') or '-')}</code></td>"
                f"<td>{html.escape(step.get('meaning') or '-')}</td>"
                f"<td>{html.escape(step.get('stopReason') or '-')}</td>"
                "</tr>"
            )
        trace_sections.append(
            f"<h2>{html.escape(row['writerVaHex'])}</h2>"
            f"<p>{html.escape(row.get('writerAlgorithm') or row.get('opcode12Algorithm') or '')}</p>"
            f"<p>{html.escape(row.get('opcode10HelperRule') or '')}</p>"
            "<h3>Activation Context</h3>"
            "<table><thead><tr><th>va</th><th>value</th><th>opcode</th><th>meaning</th><th>pointer/cns</th></tr></thead>"
            f"<tbody>{''.join(activation_rows) or '<tr><td colspan=\"5\">No activation context.</td></tr>'}</tbody></table>"
            "<h3>Trace</h3>"
            "<table><thead><tr><th>step</th><th>va</th><th>value</th><th>opcode</th><th>handler</th><th>meaning</th><th>stop</th></tr></thead>"
            f"<tbody>{''.join(trace_body)}</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 Current Writer Paths</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 Current Writer Paths</h1>",
        "  <p>Local traces for the current selector-root writers that appear before the <code>map1_01a -> map2_02d</code> frontier reader.</p>",
        f"  <p><b>Evidence refs:</b> {len(EVIDENCE_REFS)}.</p>",
        "  <p><b>promotion status:</b> <code>blocked-current-internal-only</code>.</p>",
        f"  <p><b>proof found:</b> False; <b>failed current-writer gates:</b> <code>{html.escape(','.join(FAILED_CURRENT_WRITER_PATH_GATE_IDS))}</code>; <b>missing evidence:</b> {len(CURRENT_WRITER_PATH_MISSING_EVIDENCE)}.</p>",
        f"  <h2>Evidence Refs</h2><ul>{evidence_refs}</ul>",
        "  <h2>Missing Evidence</h2>",
        "  <ul>",
        "\n".join(f"    <li>{html.escape(item)}</li>" for item in CURRENT_WRITER_PATH_MISSING_EVIDENCE),
        "  </ul>",
        "  <h2>Writer Summary</h2>",
        "  <table><thead><tr><th>writer</th><th>stream start</th><th>activation</th><th>root</th><th>maps</th><th>setup</th><th>following</th><th>next question</th></tr></thead>",
        f"  <tbody>{''.join(summary_rows) or '<tr><td colspan=\"8\">No writer rows.</td></tr>'}</tbody></table>",
        "\n".join(trace_sections),
        "</body>",
        "</html>",
        "",
    ])


def write_outputs(rows: list[dict], out_dir: Path = OUT, html_out: Path | None = None) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_current_writer_paths.json").write_text(
        json.dumps(rows, 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(rows), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--exe", type=Path, default=ROOT / "Hwanse2.exe")
    parser.add_argument("--selection-writers", type=Path, default=OUT / "save_selector_selection_writers.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()
    rows = build_rows(
        args.exe.read_bytes(),
        json.loads(args.selection_writers.read_text(encoding="utf-8")),
    )
    write_outputs(rows, args.out_dir, args.html_out)
    print(f"wrote {len(rows)} current writer path rows -> {args.out_dir / 'save_selector_current_writer_paths.json'}")


if __name__ == "__main__":
    main()