#!/usr/bin/env python3
"""Find raw event/object VM stream candidates for branch-state writer handlers."""
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


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
EVENT_HANDLER_TABLE_VA = 0x0047F1D8
EVENT_DISPATCHER_VA = 0x0041B687
EVENT_DISPATCH_CALL_VA = 0x0041B6CD
BRANCH_STATE_HANDLER_INDEXES = {
    0x25: "object/stat comparison group",
    0x26: "party/object condition group",
    0x27: "six-slot object condition group",
    0x31: "list selection initialization group",
}


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


def section_for_va(sections: list[dict], va: int) -> dict | None:
    for section in sections:
        if section["va"] <= va < section["va"] + section["raw_size"]:
            return section
    return None


def section_for_offset(sections: list[dict], offset: int) -> dict | None:
    for section in sections:
        if section["raw"] <= offset < section["raw"] + section["raw_size"]:
            return section
    return None


def dword_at_offset(exe: bytes, offset: int) -> int | None:
    if offset < 0 or offset + 4 > len(exe):
        return None
    return struct.unpack_from("<I", exe, offset)[0]


def va_for_offset(sections: list[dict], offset: int) -> int | None:
    section = section_for_offset(sections, offset)
    if section is None:
        return None
    return section["va"] + offset - section["raw"]


def pointer_target(exe: bytes, sections: list[dict], value: int) -> dict | None:
    section = section_for_va(sections, value)
    if section is None:
        return None
    return {"targetVaHex": hex32(value), "targetSection": section["name"]}


def linked_cns_near(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    offset: int,
    radius: int = 0x80,
) -> list[dict]:
    start = max(0, offset - radius)
    end = min(len(exe) - 4, offset + radius)
    rows = []
    seen = set()
    for off in range(start, end + 1, 4):
        value = dword_at_offset(exe, off)
        if value not in strings:
            continue
        key = (off, value)
        if key in seen:
            continue
        seen.add(key)
        rows.append({
            "refVaHex": hex32(va_for_offset(sections, off) or 0),
            "name": strings[value],
            "distance": off - offset,
        })
    return rows


def pointer_ref_count(exe: bytes, sections: list[dict], target: int) -> int:
    needle = struct.pack("<I", target)
    count = 0
    search = 0
    while True:
        hit = exe.find(needle, search)
        if hit < 0:
            return count
        search = hit + 1
        if section_for_offset(sections, hit) is not None:
            count += 1


def classify_candidate(exe: bytes, sections: list[dict], offset: int, index: int) -> dict:
    va = va_for_offset(sections, offset)
    aligned_off = offset - (offset % 4)
    aligned_value = dword_at_offset(exe, aligned_off)
    current_value = dword_at_offset(exe, offset)
    preceding = exe[offset - 1] if offset > 0 else None
    following = exe[offset + 2] if offset + 2 < len(exe) else None
    aligned_pointer = pointer_target(exe, sections, aligned_value) if aligned_value is not None else None
    current_pointer = pointer_target(exe, sections, current_value) if current_value is not None else None
    if offset % 4 == 0 and current_pointer:
        classification = "aligned pointer bytes, likely data not command start"
        confidence = "low"
    elif preceding == 0x10:
        classification = "embedded after 0x10 prefix; possible operand overlap, needs VM length proof"
        confidence = "weak"
    elif following in {0x00, 0x01, 0x02, 0x03, 0x04}:
        classification = "raw 0x40-index byte pair in data; possible event/object command start"
        confidence = "medium"
    else:
        classification = "raw byte overlap in data; unproven command start"
        confidence = "low"
    return {
        "vaHex": hex32(va or 0),
        "fileOffsetHex": f"0x{offset:06x}",
        "indexHex": f"0x{index:02x}",
        "handlerLabel": BRANCH_STATE_HANDLER_INDEXES[index],
        "offsetMod4": offset % 4,
        "precedingByteHex": f"0x{preceding:02x}" if preceding is not None else None,
        "followingByteHex": f"0x{following:02x}" if following is not None else None,
        "alignedDwordVaHex": hex32(va_for_offset(sections, aligned_off) or 0),
        "alignedDwordHex": hex32(aligned_value) if aligned_value is not None else None,
        "alignedDwordPointer": aligned_pointer,
        "currentDwordHex": hex32(current_value) if current_value is not None else None,
        "currentDwordPointer": current_pointer,
        "pointerRefCount": pointer_ref_count(exe, sections, va or 0) if va is not None else 0,
        "classification": classification,
        "confidence": confidence,
    }


def build_summary(exe: bytes) -> dict:
    sections = read_sections(exe)
    strings = find_cns_strings(exe, sections)
    candidates = []
    for section in sections:
        if section["name"] not in {".data", ".rdata"}:
            continue
        raw_start = section["raw"]
        raw = exe[raw_start: raw_start + section["raw_size"]]
        for index in BRANCH_STATE_HANDLER_INDEXES:
            pattern = bytes([0x40, index])
            search = 0
            while True:
                hit = raw.find(pattern, search)
                if hit < 0:
                    break
                search = hit + 1
                offset = raw_start + hit
                row = classify_candidate(exe, sections, offset, index)
                row["section"] = section["name"]
                row["nearbyCns"] = linked_cns_near(exe, sections, strings, offset)
                candidates.append(row)
    candidates.sort(key=lambda row: (row["indexHex"], row["vaHex"]))
    by_index = {}
    by_confidence = {}
    for row in candidates:
        by_index[row["indexHex"]] = by_index.get(row["indexHex"], 0) + 1
        by_confidence[row["confidence"]] = by_confidence.get(row["confidence"], 0) + 1
    return {
        "scope": "raw .data/.rdata byte pairs that could dispatch branch-state event/object handlers",
        "eventObjectDispatcher": {
            "dispatcherVaHex": hex32(EVENT_DISPATCHER_VA),
            "dispatchCallVaHex": hex32(EVENT_DISPATCH_CALL_VA),
            "handlerTableVaHex": hex32(EVENT_HANDLER_TABLE_VA),
            "dispatchRule": "stream byte0 must be 0x40; stream byte1 indexes dword [byte1*4 + 0x0047f1d8]",
        },
        "candidateCount": len(candidates),
        "candidateCountsByIndex": by_index,
        "candidateCountsByConfidence": by_confidence,
        "conclusion": (
            "Raw 0x40-index byte pairs exist for branch-state writer handlers, but many are embedded in pointers or operands. "
            "This report separates byte overlap from command-start evidence; none of these rows alone proves that selector root 2:0 "
            "executes a branch-state writer before the current frontier."
        ),
        "candidates": candidates,
    }


def markdown(summary: dict) -> str:
    lines = [
        "# Event/Object Branch State Stream Candidates",
        "",
        "Raw `0x40 xx` byte-pair candidates for event/object VM handlers that write `primaryBranchState`.",
        "",
        f"- dispatcher: `{summary['eventObjectDispatcher']['dispatcherVaHex']}`",
        f"- dispatch call: `{summary['eventObjectDispatcher']['dispatchCallVaHex']}`",
        f"- handler table: `{summary['eventObjectDispatcher']['handlerTableVaHex']}`",
        f"- candidates: {summary['candidateCount']}",
        "",
        summary["conclusion"],
        "",
        "| va | index | confidence | classification | aligned dword | nearby CNS |",
        "| --- | --- | --- | --- | --- | --- |",
    ]
    for row in summary["candidates"]:
        cns = ", ".join(item["name"] for item in row.get("nearbyCns") or []) or "-"
        lines.append(
            f"| `{row['vaHex']}` | `{row['indexHex']}` {row['handlerLabel']} | {row['confidence']} | "
            f"{row['classification']} | `{row['alignedDwordHex']}` | {cns} |"
        )
    lines.append("")
    return "\n".join(lines)


def html_page(summary: dict) -> str:
    rows = []
    for row in summary["candidates"]:
        cns = "<br>".join(html.escape(item["name"]) for item in row.get("nearbyCns") or []) or "-"
        rows.append(
            "<tr>"
            f"<td><code>{html.escape(row['vaHex'])}</code></td>"
            f"<td><code>{html.escape(row['indexHex'])}</code><br>{html.escape(row['handlerLabel'])}</td>"
            f"<td>{html.escape(row['confidence'])}</td>"
            f"<td>{html.escape(row['classification'])}</td>"
            f"<td><code>{html.escape(row.get('alignedDwordHex') or '-')}</code></td>"
            f"<td>{cns}</td>"
            "</tr>"
        )
    return "\n".join([
        "<!doctype html>",
        '<html lang="en">',
        "<head>",
        '  <meta charset="utf-8">',
        '  <meta name="viewport" content="width=device-width, initial-scale=1">',
        "  <title>Event/Object Branch State Stream Candidates</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>Event/Object Branch State Stream Candidates</h1>",
        f"  <p>Dispatcher: <code>{html.escape(summary['eventObjectDispatcher']['dispatcherVaHex'])}</code>, "
        f"call <code>{html.escape(summary['eventObjectDispatcher']['dispatchCallVaHex'])}</code>, "
        f"table <code>{html.escape(summary['eventObjectDispatcher']['handlerTableVaHex'])}</code>.</p>",
        f"  <p>{html.escape(summary['conclusion'])}</p>",
        "  <table><thead><tr><th>va</th><th>index</th><th>confidence</th><th>classification</th><th>aligned dword</th><th>nearby CNS</th></tr></thead>",
        f"  <tbody>{''.join(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 / "event_object_branch_state_stream_candidates.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("--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 event/object branch-state stream candidates -> {args.out_dir / 'event_object_branch_state_stream_candidates.json'}")


if __name__ == "__main__":
    main()