#!/usr/bin/env python3
"""Trace save-selector leaf streams using script handler advance candidates."""
from __future__ import annotations

import argparse
import html
import json
import struct
from pathlib import Path

from probe_exe_scene_tables import find_cns_strings, read_sections, va_to_offset
from summarize_script_handler_table import DEFAULT_HANDLER_VA, HANDLER_TABLE_VA
from summarize_script_handler_table import analyze_stream_effect, dword_at, section_name_for_va


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


def byte_at(exe: bytes, sections: list[dict], va: int) -> int | None:
    offset = va_to_offset(sections, va)
    if offset is None or offset >= len(exe):
        return None
    return exe[offset]


def u32_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 handler_entry(exe: bytes, sections: list[dict], opcode: int) -> dict:
    entry_va = HANDLER_TABLE_VA + opcode * 4
    handler_va = dword_at(exe, sections, entry_va)
    section = section_name_for_va(sections, handler_va) if handler_va is not None else None
    effect = {}
    if handler_va == DEFAULT_HANDLER_VA:
        effect = {"fixedAdvances": [{"bytes": 4}], "streamReads": [], "streamStores": []}
    elif handler_va is not None and section == ".text":
        effect = analyze_stream_effect(exe, sections, handler_va)
    advances = sorted({item.get("bytes") for item in effect.get("fixedAdvances", []) if item.get("bytes")})
    return {
        "opcode": opcode,
        "opcodeHex": f"0x{opcode:02x}",
        "entryVaHex": f"0x{entry_va:08x}",
        "handlerVa": handler_va,
        "handlerVaHex": f"0x{handler_va:08x}" if handler_va is not None else None,
        "handlerSection": section,
        "isDefaultHandler": handler_va == DEFAULT_HANDLER_VA,
        "fixedAdvances": advances,
        "canJumpToDwordAtPlus4": effect.get("canJumpToDwordAtPlus4") is True,
        "writesStream": bool(effect.get("streamStores")),
        "streamReads": effect.get("streamReads") or [],
    }


def branch_note(handler_va: int | None) -> str | None:
    if handler_va == 0x0040B4E6:
        return (
            "if table[(context+0xa8)[stream+2]] == 1 then fall through +8; "
            "table is 0x59e370 when stream+1 is 0, otherwise 0x59e360; "
            "else jump to dword [stream+4]"
        )
    if handler_va == 0x0040FA6D:
        return "jump to dword [stream+4] when encounter-state flags are set; otherwise fall through +8"
    if handler_va == 0x0040F426:
        return "condition helper can jump to dword [stream+4] or fall through +8"
    return None


def trace_stream(exe: bytes, sections: list[dict], strings: dict[int, str], stream_va: int, max_steps: int) -> list[dict]:
    rows = []
    seen = set()
    va = stream_va
    for step in range(max_steps):
        if va in seen:
            rows.append({"step": step, "vaHex": f"0x{va:08x}", "stopReason": "loop"})
            break
        seen.add(va)
        opcode = byte_at(exe, sections, va)
        value = u32_at(exe, sections, va)
        if opcode is None or value is None:
            rows.append({"step": step, "vaHex": f"0x{va:08x}", "stopReason": "unreadable"})
            break
        handler = handler_entry(exe, sections, opcode)
        row = {
            "step": step,
            "va": va,
            "vaHex": f"0x{va:08x}",
            "value": value,
            "valueHex": f"0x{value:08x}",
            **handler,
        }
        if value in strings:
            row["cns"] = strings[value]
        elif va_to_offset(sections, value) is not None:
            row["pointer"] = True
            row["pointerHex"] = f"0x{value:08x}"
        if handler["canJumpToDwordAtPlus4"]:
            target = u32_at(exe, sections, va + 4)
            row["branchTargetHex"] = f"0x{target:08x}" if target is not None else None
            row["fallthroughVaHex"] = f"0x{va + 8:08x}"
            note = branch_note(handler.get("handlerVa"))
            if note:
                row["branchNote"] = note
        rows.append(row)
        advances = handler["fixedAdvances"]
        if len(advances) == 1 and not handler["canJumpToDwordAtPlus4"]:
            va += advances[0]
            continue
        if len(advances) == 1 and handler["canJumpToDwordAtPlus4"]:
            row["stopReason"] = "branch-or-fallthrough"
            break
        if not advances:
            row["stopReason"] = "no-fixed-advance"
            break
        row["stopReason"] = "multiple-advances"
        break
    return rows


def build_rows(exe: bytes, leaf_rows: list[dict], max_steps: int = 32) -> list[dict]:
    sections = read_sections(exe)
    strings = find_cns_strings(exe, sections)
    rows = []
    for leaf in leaf_rows:
        for kind, key in [("leaf", "leafPointerHex"), ("nested", "nestedPointerHex")]:
            stream_hex = leaf.get(key)
            if not stream_hex:
                continue
            stream_va = int(stream_hex, 16)
            rows.append({
                "source": leaf.get("source"),
                "target": leaf.get("target"),
                "leafPointerHex": leaf.get("leafPointerHex"),
                "streamKind": kind,
                "streamVaHex": stream_hex,
                "trace": trace_stream(exe, sections, strings, stream_va, max_steps),
            })
    return rows


def markdown(rows: list[dict]) -> str:
    lines = [
        "# Save Selector Stream Traces",
        "",
        "Linear bytecode traces for save-selector leaf streams. These traces follow only handlers with one fixed stream advance.",
        "Rows with branch targets, missing advances, or multiple possible advances are stopped rather than guessed through.",
        "",
        f"Traced streams: {len(rows)}.",
        "",
    ]
    for row in rows:
        lines.extend([
            f"## {row['streamVaHex']} ({row['streamKind']}, leaf {row['leafPointerHex']})",
            "",
            "| step | va | value | opcode | handler | effect | branch target | fallthrough | stop | note |",
            "| ---: | --- | --- | --- | --- | --- | --- | --- | --- | --- |",
        ])
        for item in row["trace"]:
            advances = "/".join(f"+{value}" for value in item.get("fixedAdvances", [])) or "-"
            effect_bits = [f"advance {advances}" if advances != "-" else "-"]
            if item.get("writesStream"):
                effect_bits.append("writes stream")
            if item.get("canJumpToDwordAtPlus4"):
                effect_bits.append("branch")
            lines.append(
                f"| {item['step']} | `{item['vaHex']}` | `{item.get('valueHex', '-')}` | "
                f"`{item.get('opcodeHex', '-')}` | `{item.get('handlerVaHex', '-')}` | "
                f"{'; '.join(effect_bits)} | `{item.get('branchTargetHex') or '-'}` | "
                f"`{item.get('fallthroughVaHex') or '-'}` | {item.get('stopReason') or '-'} | "
                f"{item.get('branchNote') or '-'} |"
            )
        lines.append("")
    return "\n".join(lines)


def html_page(rows: list[dict]) -> str:
    sections = []
    for row in rows:
        body = []
        for item in row["trace"]:
            advances = "/".join(f"+{value}" for value in item.get("fixedAdvances", [])) or "-"
            effect = [f"advance {advances}" if advances != "-" else "-"]
            if item.get("writesStream"):
                effect.append("writes stream")
            if item.get("canJumpToDwordAtPlus4"):
                effect.append("branch")
            body.append(
                "<tr>"
                f"<td>{item['step']}</td>"
                f"<td><code>{html.escape(item['vaHex'])}</code></td>"
                f"<td><code>{html.escape(item.get('valueHex', '-'))}</code></td>"
                f"<td><code>{html.escape(item.get('opcodeHex', '-'))}</code></td>"
                f"<td><code>{html.escape(item.get('handlerVaHex') or '-')}</code></td>"
                f"<td>{html.escape('; '.join(effect))}</td>"
                f"<td><code>{html.escape(item.get('branchTargetHex') or '-')}</code></td>"
                f"<td><code>{html.escape(item.get('fallthroughVaHex') or '-')}</code></td>"
                f"<td>{html.escape(item.get('stopReason') or '-')}</td>"
                f"<td>{html.escape(item.get('branchNote') or '-')}</td>"
                "</tr>"
            )
        sections.append(
            f"<h2>{html.escape(row['streamVaHex'])} ({html.escape(row['streamKind'])}, leaf {html.escape(row['leafPointerHex'])})</h2>"
            "<table><thead><tr><th>step</th><th>va</th><th>value</th><th>opcode</th><th>handler</th><th>effect</th><th>branch target</th><th>fallthrough</th><th>stop</th><th>note</th></tr></thead>"
            f"<tbody>{''.join(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 Stream Traces</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; position: sticky; top: 0; }",
        "    code { color: #f5d76e; }",
        "  </style>",
        "</head>",
        "<body>",
        "  <h1>Save Selector Stream Traces</h1>",
        "  <p>Linear bytecode traces for save-selector leaf streams. These traces follow only handlers with one fixed stream advance.</p>",
        "  <p>Rows with branch targets, missing advances, or multiple possible advances are stopped rather than guessed through.</p>",
        f"  <p>Traced streams: {len(rows)}.</p>",
        "\n".join(sections),
        "</body>",
        "</html>",
        "",
    ])


def write_outputs(rows: list[dict], out_dir: Path = OUT) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "save_selector_stream_traces.json").write_text(
        json.dumps(rows, 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("--leaf-streams", type=Path, default=OUT / "save_selector_leaf_streams.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    parser.add_argument("--max-steps", type=int, default=32)
    args = parser.parse_args()
    rows = build_rows(
        args.exe.read_bytes(),
        json.loads(args.leaf_streams.read_text(encoding="utf-8")),
        args.max_steps,
    )
    write_outputs(rows, args.out_dir)
    print(f"wrote {len(rows)} save selector stream traces -> {args.out_dir / 'save_selector_stream_traces.json'}")


if __name__ == "__main__":
    main()