#!/usr/bin/env python3
"""Probe scene manifest records for nearby prompt/text references."""
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"


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


def load_json(path: Path):
    return json.loads(path.read_text(encoding="utf-8"))


def prompt_targets(prompts_summary: dict) -> dict[int, list[dict]]:
    targets: dict[int, list[dict]] = {}
    for prompt in prompts_summary.get("prompts") or []:
        for key in ("renderVaHex", "startVaHex", "endVaHex", "waitVaHex"):
            value = prompt.get(key)
            if not isinstance(value, str) or not value.startswith("0x"):
                continue
            targets.setdefault(int(value, 16), []).append(
                {
                    "promptId": prompt.get("id", ""),
                    "targetKind": key,
                    "sample": (prompt.get("displayText") or prompt.get("text") or "").replace("\n", " / ")[:120],
                }
            )
        for value in prompt.get("lineVas") or []:
            if not isinstance(value, str) or not value.startswith("0x"):
                continue
            targets.setdefault(int(value, 16), []).append(
                {
                    "promptId": prompt.get("id", ""),
                    "targetKind": "lineVa",
                    "sample": (prompt.get("displayText") or prompt.get("text") or "").replace("\n", " / ")[:120],
                }
            )
    return targets


def text_targets(korean_summary: dict) -> dict[int, dict]:
    targets: dict[int, dict] = {}
    for row in korean_summary.get("candidates") or []:
        if isinstance(row.get("va"), int):
            targets[row["va"]] = {
                "targetKind": "koreanText",
                "sample": row.get("text", "")[:120],
            }
    for cluster in korean_summary.get("refClusters") or []:
        for entry in cluster.get("entries") or []:
            text_va = int(entry["textVaHex"], 16)
            ref_va = int(entry["refVaHex"], 16)
            targets[text_va] = {
                "targetKind": "koreanText",
                "sample": entry.get("text", "")[:120],
            }
            targets[ref_va] = {
                "targetKind": "koreanTextTableEntry",
                "sample": entry.get("text", "")[:120],
            }
    return targets


def dwords_in_window(exe: bytes, sections: list[dict], center_va: int, radius: int) -> list[tuple[int, int]]:
    start_va = center_va - radius
    end_va = center_va + radius
    va = start_va - (start_va % 4)
    rows: list[tuple[int, int]] = []
    while va <= end_va:
        offset = va_to_offset(sections, va)
        if offset is not None and 0 <= offset <= len(exe) - 4:
            rows.append((va, struct.unpack_from("<I", exe, offset)[0]))
        va += 4
    return rows


def scan_radius(
    exe: bytes,
    sections: list[dict],
    manifest: list[dict],
    prompt_index: dict[int, list[dict]],
    text_index: dict[int, dict],
    radius: int,
) -> dict:
    prompt_hits: list[dict] = []
    text_hits: list[dict] = []
    for record in manifest:
        record_va = record["recordVa"]
        for ref_va, value in dwords_in_window(exe, sections, record_va, radius):
            for target in prompt_index.get(value) or []:
                prompt_hits.append(
                    {
                        "map": record.get("map", ""),
                        "sceneIdHex": record.get("sceneIdHex", ""),
                        "recordVaHex": record.get("recordVaHex") or hex32(record_va),
                        "refVaHex": hex32(ref_va),
                        "valueHex": hex32(value),
                        "distance": abs(ref_va - record_va),
                        **target,
                    }
                )
            text_target = text_index.get(value)
            if text_target:
                text_hits.append(
                    {
                        "map": record.get("map", ""),
                        "sceneIdHex": record.get("sceneIdHex", ""),
                        "recordVaHex": record.get("recordVaHex") or hex32(record_va),
                        "refVaHex": hex32(ref_va),
                        "valueHex": hex32(value),
                        "distance": abs(ref_va - record_va),
                        **text_target,
                    }
                )
    return {
        "radius": radius,
        "radiusHex": f"0x{radius:x}",
        "promptRefCount": len(prompt_hits),
        "textRefCount": len(text_hits),
        "promptRefs": sorted(prompt_hits, key=lambda row: (row["distance"], row["map"], row["refVaHex"]))[:80],
        "textRefs": sorted(text_hits, key=lambda row: (row["distance"], row["map"], row["refVaHex"]))[:80],
    }


def build_summary(exe: bytes, manifest: list[dict], prompts: dict, korean_text: dict) -> dict:
    sections = read_sections(exe)
    prompt_index = prompt_targets(prompts)
    text_index = text_targets(korean_text)
    scans = [
        scan_radius(exe, sections, manifest, prompt_index, text_index, radius)
        for radius in (0x100, 0x400, 0x1000)
    ]
    strict = next(scan for scan in scans if scan["radius"] == 0x400)
    conclusion = (
        "No direct story prompt or Korean text pointer was found within 0x400 bytes of scene manifest records. "
        "Wider 0x1000 hits exist, but they overlap unrelated global tables and should remain audit-only candidates."
        if strict["promptRefCount"] == 0 and strict["textRefCount"] == 0
        else "At least one direct prompt/text pointer was found within the strict scene-record window and needs opcode tracing."
    )
    return {
        "scope": "Scene manifest record windows checked for direct story prompt/text address references.",
        "sceneRecordCount": len(manifest),
        "promptTargetAddressCount": len(prompt_index),
        "koreanTextTargetAddressCount": len(text_index),
        "strictEvidenceRadiusHex": "0x400",
        "conclusion": conclusion,
        "scans": scans,
    }


def html_page(summary: dict) -> str:
    rows = []
    for scan in summary["scans"]:
        rows.append(
            "<tr>"
            f"<td><code>{html.escape(scan['radiusHex'])}</code></td>"
            f"<td>{scan['promptRefCount']}</td>"
            f"<td>{scan['textRefCount']}</td>"
            "</tr>"
        )
    samples = []
    for row in summary["scans"][-1].get("promptRefs") or []:
        samples.append(
            "<tr>"
            f"<td>{html.escape(row['map'])}</td>"
            f"<td>{html.escape(row['sceneIdHex'])}</td>"
            f"<td><code>{html.escape(row['recordVaHex'])}</code></td>"
            f"<td><code>{html.escape(row['refVaHex'])}</code></td>"
            f"<td><code>{html.escape(row['valueHex'])}</code></td>"
            f"<td>{html.escape(row['promptId'])} / {html.escape(row['targetKind'])}</td>"
            f"<td>{row['distance']}</td>"
            f"<td>{html.escape(row['sample'])}</td>"
            "</tr>"
        )
    return "\n".join(
        [
            "<!doctype html>",
            '<html lang="ko">',
            "<head>",
            '  <meta charset="utf-8">',
            '  <meta name="viewport" content="width=device-width, initial-scale=1">',
            "  <title>Scene Prompt Reference Probe</title>",
            "  <style>body{font-family:system-ui,sans-serif;background:#101010;color:#eee;margin:24px}table{border-collapse:collapse;width:100%;margin:16px 0}td,th{border:1px solid #333;padding:6px 8px;vertical-align:top}th{background:#1d1d1d}code{color:#f5d76e}.muted{color:#aaa}</style>",
            "</head>",
            "<body>",
            "  <h1>Scene Prompt Reference Probe</h1>",
            f"  <p>{html.escape(summary['conclusion'])}</p>",
            f"  <p class=\"muted\">scene records={summary['sceneRecordCount']} · prompt targets={summary['promptTargetAddressCount']} · text targets={summary['koreanTextTargetAddressCount']}</p>",
            "  <table><thead><tr><th>radius</th><th>prompt refs</th><th>text refs</th></tr></thead>",
            f"  <tbody>{''.join(rows)}</tbody></table>",
            "  <h2>Wide Audit Prompt Samples</h2>",
            "  <table><thead><tr><th>map</th><th>scene</th><th>record</th><th>ref</th><th>value</th><th>target</th><th>distance</th><th>sample</th></tr></thead>",
            f"  <tbody>{''.join(samples) or '<tr><td colspan=\"8\">none</td></tr>'}</tbody></table>",
            "</body>",
            "</html>",
            "",
        ]
    )


def write_outputs(summary: dict, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "scene_prompt_reference_probe.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("--scene-manifest", type=Path, default=OUT / "scene_manifest.json")
    parser.add_argument("--story-prompts", type=Path, default=OUT / "story_prompts.json")
    parser.add_argument("--korean-text", type=Path, default=OUT / "korean_text_candidates.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()

    summary = build_summary(
        args.exe.read_bytes(),
        load_json(args.scene_manifest),
        load_json(args.story_prompts),
        load_json(args.korean_text),
    )
    write_outputs(summary, args.out_dir)
    strict = next(scan for scan in summary["scans"] if scan["radius"] == 0x400)
    print(
        "wrote scene prompt reference probe -> "
        f"{args.out_dir / 'scene_prompt_reference_probe.json'} "
        f"(strict prompt={strict['promptRefCount']}, text={strict['textRefCount']})"
    )


if __name__ == "__main__":
    main()