#!/usr/bin/env python3
"""Scan map1_01a/map2_02d resource references for nearby point tables."""
from __future__ import annotations

import argparse
import html
import json
import struct
import sys
from collections import Counter
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent))

from extract_scene_events import read_point_table
from probe_exe_scene_tables import find_cns_strings, offset_to_va, read_sections, va_to_offset


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
SOURCE = "map1_01a"
TARGET = "map2_02d"
WINDOW_BEFORE = 0x180
WINDOW_AFTER = 0x220
POINT_TABLE_SCAN_SPAN = 0x80
REPORT_POINT_LIMIT = 12


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


def load_json(path: Path, fallback: Any) -> Any:
    if not path.exists():
        return fallback
    return json.loads(path.read_text(encoding="utf-8"))


def parse_maps_js(path: Path) -> dict:
    text = path.read_text(encoding="utf-8")
    prefix = "window.HWANSE_MAPS = "
    if not text.startswith(prefix):
        raise ValueError(f"{path} does not contain the expected maps.js wrapper")
    return json.loads(text[len(prefix):].rstrip(";\n"))


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 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 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 dword_at_va(exe: bytes, sections: list[dict], va: int) -> int | None:
    offset = va_to_offset(sections, va)
    if offset is None:
        return None
    return dword_at_offset(exe, offset)


def is_scene_record(exe: bytes, strings: dict[int, str], offset: int) -> bool:
    value = dword_at_offset(exe, offset)
    scene_id = dword_at_offset(exe, offset + 4)
    zero = dword_at_offset(exe, offset + 8)
    name = strings.get(value or -1)
    return bool(
        name
        and name.startswith("map")
        and name.endswith(".cns")
        and isinstance(scene_id, int)
        and 0 < scene_id < 0x100000
        and scene_id & 0xFF in {0x00, 0x18}
        and zero == 0
    )


def cns_refs_in_window(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    start_va: int,
    end_va: int,
) -> list[dict]:
    rows = []
    start = va_to_offset(sections, start_va)
    end = va_to_offset(sections, end_va - 1)
    if start is None or end is None:
        return rows
    end += 1
    aligned = start - (start % 4)
    for offset in range(aligned, min(end, len(exe) - 4), 4):
        value = dword_at_offset(exe, offset)
        name = strings.get(value or -1)
        if not name:
            continue
        va = offset_to_va(sections, offset)
        rows.append({
            "refVa": va,
            "refVaHex": hex32(va or 0),
            "filename": name,
            "mapStem": name[:-4] if name.endswith(".cns") else name,
        })
    return rows


def pointer_ref_count(exe: bytes, sections: list[dict], value: int, sample_limit: int = 5) -> dict:
    needle = struct.pack("<I", value)
    count = 0
    samples = []
    search = 0
    while True:
        hit = exe.find(needle, search)
        if hit < 0:
            break
        search = hit + 1
        section = section_for_offset(sections, hit)
        if section is None:
            continue
        ref_va = section["va"] + hit - section["raw"]
        count += 1
        if len(samples) < sample_limit:
            samples.append({
                "section": section["name"],
                "refVaHex": hex32(ref_va),
            })
    return {"count": count, "samples": samples}


def route_exit_samples(map_exit_candidates: list[dict]) -> list[dict]:
    rows = []
    for map_row in map_exit_candidates:
        if map_row.get("map") != SOURCE:
            continue
        if TARGET not in (map_row.get("selectorTargetCandidates") or []):
            continue
        for candidate in map_row.get("exitCandidates") or []:
            sample = candidate.get("sample") or {}
            x = sample.get("x")
            y = sample.get("y")
            if isinstance(x, int) and isinstance(y, int):
                rows.append({
                    "side": candidate.get("side"),
                    "x": x,
                    "y": y,
                    "autoTrigger": candidate.get("autoTrigger"),
                    "edgeDistance": candidate.get("edgeDistance"),
                })
    unique = {}
    for row in rows:
        unique[(row["side"], row["x"], row["y"])] = row
    return list(unique.values())


def nearest_scene_record_before(exe: bytes, sections: list[dict], strings: dict[int, str], va: int) -> dict | None:
    section = section_for_va(sections, va)
    if section is None:
        return None
    start = max(section["va"], va - 0x300)
    for candidate_va in range(va - 4, start - 1, -4):
        offset = va_to_offset(sections, candidate_va)
        if offset is None or not is_scene_record(exe, strings, offset):
            continue
        value = dword_at_offset(exe, offset)
        scene_id = dword_at_offset(exe, offset + 4)
        name = strings.get(value or -1)
        return {
            "recordVa": candidate_va,
            "recordVaHex": hex32(candidate_va),
            "filename": name,
            "mapStem": name[:-4] if isinstance(name, str) and name.endswith(".cns") else name,
            "sceneIdHex": f"0x{scene_id:04x}" if isinstance(scene_id, int) else None,
        }
    return None


def classify_point_candidate(row: dict) -> str:
    if row.get("exactRouteExitHitCount"):
        return "route-exit-point-table-candidate"
    raw_count = int(row.get("rawPointCount") or 0)
    in_bounds_count = int(row.get("inBoundsPointCount") or 0)
    first = row.get("firstRawPoints") or []
    first_point = first[0] if first else {}
    x = first_point.get("x")
    y = first_point.get("y")
    owner = row.get("ownerSceneRecord") or {}
    if raw_count == 1:
        if x == 63 and y == 0:
            return "sentinel-0x003f-singleton"
        if y == 0 and isinstance(x, int) and x >= 128 and x % 16 == 0:
            return "resource-size-singleton"
        if in_bounds_count == 0:
            return "out-of-bounds-singleton"
        return "single-in-bounds-point"
    if raw_count <= 3:
        if in_bounds_count == raw_count:
            return "short-in-bounds-scalar-run"
        return "short-out-of-bounds-scalar-run"
    if owner and in_bounds_count == raw_count:
        return "owned-point-table-non-route"
    if x == 33 and y == 1:
        return "script-command-shaped-run"
    if in_bounds_count == 0:
        return "out-of-bounds-run"
    return "point-table-near-resource-reference"


def point_candidates_for_ref(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    map_data: dict,
    ref: dict,
    exit_points: set[tuple[int, int]],
) -> list[dict]:
    source_map = map_data.get(SOURCE) or {}
    width = int(source_map.get("width") or 0)
    height = int(source_map.get("height") or 0)
    if not width or not height:
        return []
    ref_va = ref["refVa"]
    start_va = max(ref["sectionVa"], ref_va - WINDOW_BEFORE)
    end_va = min(ref["sectionEndVa"], ref_va + WINDOW_AFTER)
    rows = []
    for field_va in range(start_va - (start_va % 4), end_va, 4):
        value = dword_at_va(exe, sections, field_va)
        if value is None or value in strings or not (0x00400000 <= value < 0x00600000):
            continue
        if va_to_offset(sections, value) is None:
            continue
        raw_points, in_bounds = read_point_table(exe, sections, value, width, height)
        if not raw_points:
            continue
        exact_hits = [
            point
            for point in raw_points
            if (point["x"], point["y"]) in exit_points
        ]
        refs = pointer_ref_count(exe, sections, value)
        owner = nearest_scene_record_before(exe, sections, strings, field_va)
        row = {
            "fieldVa": field_va,
            "fieldVaHex": hex32(field_va),
            "payloadVa": value,
            "payloadVaHex": hex32(value),
            "ownerSceneRecord": owner,
            "distanceFromResourceRef": field_va - ref_va,
            "rawPointCount": len(raw_points),
            "inBoundsPointCount": len(in_bounds),
            "firstRawPoints": raw_points[:REPORT_POINT_LIMIT],
            "firstInBoundsPoints": in_bounds[:REPORT_POINT_LIMIT],
            "exactRouteExitHitCount": len(exact_hits),
            "exactRouteExitHits": exact_hits[:REPORT_POINT_LIMIT],
            "pointerRefCount": refs["count"],
            "pointerRefs": refs["samples"],
            "scanSpanHex": f"{hex32(value)}..{hex32(value + POINT_TABLE_SCAN_SPAN)}",
        }
        row["classification"] = classify_point_candidate(row)
        rows.append(row)
    return rows


def resource_refs(exe: bytes, sections: list[dict], strings: dict[int, str], filename: str) -> list[dict]:
    string_vas = [va for va, name in strings.items() if name == filename]
    rows = []
    for string_va in string_vas:
        needle = struct.pack("<I", string_va)
        search = 0
        while True:
            hit = exe.find(needle, search)
            if hit < 0:
                break
            search = hit + 1
            section = section_for_offset(sections, hit)
            if section is None:
                continue
            ref_va = section["va"] + hit - section["raw"]
            scene_record = is_scene_record(exe, strings, hit)
            scene_id = dword_at_offset(exe, hit + 4) if scene_record else None
            rows.append({
                "filename": filename,
                "mapStem": filename[:-4],
                "stringVa": string_va,
                "stringVaHex": hex32(string_va),
                "refVa": ref_va,
                "refVaHex": hex32(ref_va),
                "refSection": section["name"],
                "sectionVa": section["va"],
                "sectionEndVa": section["va"] + section["raw_size"],
                "sceneRecordLike": scene_record,
                "sceneIdHex": f"0x{scene_id:04x}" if isinstance(scene_id, int) else None,
            })
    rows.sort(key=lambda row: (row["refVa"], row["stringVa"]))
    return rows


def build_summary(
    exe: bytes,
    map_data: dict,
    map_exit_candidates: list[dict] | None = None,
) -> dict:
    sections = read_sections(exe)
    strings = find_cns_strings(exe, sections)
    exit_samples = route_exit_samples(map_exit_candidates or [])
    exit_points = {(row["x"], row["y"]) for row in exit_samples}
    records = []
    for filename in (f"{SOURCE}.cns", f"{TARGET}.cns"):
        for ref in resource_refs(exe, sections, strings, filename):
            start_va = max(ref["sectionVa"], ref["refVa"] - WINDOW_BEFORE)
            end_va = min(ref["sectionEndVa"], ref["refVa"] + WINDOW_AFTER)
            point_rows = point_candidates_for_ref(exe, sections, strings, map_data, ref, exit_points)
            cns_rows = cns_refs_in_window(exe, sections, strings, start_va, end_va)
            linked_maps = sorted({
                row["mapStem"]
                for row in cns_rows
                if row["mapStem"].startswith("map") and "_" in row["mapStem"]
            })
            strict_candidates = [
                row for row in point_rows
                if row["exactRouteExitHitCount"] and TARGET in linked_maps
            ]
            records.append({
                **{key: value for key, value in ref.items() if not key.startswith("section")},
                "windowHex": f"{hex32(start_va)}..{hex32(end_va)}",
                "nearbyMapResources": linked_maps,
                "nearbyCnsRefCount": len(cns_rows),
                "pointCandidateCount": len(point_rows),
                "pointCandidateClassCounts": dict(Counter(
                    point["classification"] for point in point_rows
                )),
                "pointCandidates": point_rows,
                "strictSourceTargetCandidateCount": len(strict_candidates),
            })
    source_records = [row for row in records if row["mapStem"] == SOURCE]
    target_records = [row for row in records if row["mapStem"] == TARGET]
    point_count = sum(row["pointCandidateCount"] for row in records)
    exact_count = sum(
        1
        for row in records
        for point in row["pointCandidates"]
        if point["exactRouteExitHitCount"]
    )
    strict_count = sum(row["strictSourceTargetCandidateCount"] for row in records)
    current_frontier = next(
        (
            row for row in source_records
            if row["refVaHex"] == "0x00542b44"
            and TARGET in row.get("nearbyMapResources", [])
        ),
        None,
    )
    point_class_counts = dict(Counter(
        point["classification"]
        for record in records
        for point in record["pointCandidates"]
    ))
    current_frontier_point_class_counts = dict(Counter(
        point["classification"]
        for point in ((current_frontier or {}).get("pointCandidates") or [])
    ))
    current_frontier_route_exit_point_hits = sum(
        point.get("exactRouteExitHitCount", 0)
        for point in ((current_frontier or {}).get("pointCandidates") or [])
    )
    current_frontier_non_route_singleton_count = sum(
        count
        for key, count in current_frontier_point_class_counts.items()
        if key in {
            "sentinel-0x003f-singleton",
            "resource-size-singleton",
            "out-of-bounds-singleton",
        }
    )
    conclusion = (
        "The direct resource-reference scan finds the expected scene-list references for map1_01a and map2_02d, "
        "including the current selector-only frontier at 0x00542b44. Nearby point-table-like payloads still do not "
        "produce a strict map1_01a route-exit hit for map2_02d. The current frontier candidates are non-route "
        "singletons or resource-size payloads, so this scan does not promote the transition."
    )
    return {
        "source": SOURCE,
        "target": TARGET,
        "windowBeforeHex": hex32(WINDOW_BEFORE),
        "windowAfterHex": hex32(WINDOW_AFTER),
        "routeExitSamples": exit_samples,
        "sourceReferenceCount": len(source_records),
        "targetReferenceCount": len(target_records),
        "resourceReferenceCount": len(records),
        "pointCandidateCount": point_count,
        "pointCandidateClassCounts": point_class_counts,
        "routeExitPointCandidateCount": exact_count,
        "strictSourceTargetCandidateCount": strict_count,
        "currentFrontierReferenceFound": current_frontier is not None,
        "currentFrontierReference": current_frontier,
        "currentFrontierPointCandidateClassCounts": current_frontier_point_class_counts,
        "currentFrontierRouteExitPointHitCount": current_frontier_route_exit_point_hits,
        "currentFrontierNonRouteSingletonPointCandidateCount": (
            current_frontier_non_route_singleton_count
        ),
        "promotionStatus": "blocked",
        "remainingProofs": [
            "find a point/hotspot table that belongs to map1_01a and targets map2_02d",
            "or prove the selector-only resource cluster executes from a non-coordinate runtime trigger",
            "or capture a runtime trace / real selector 2:0 savedata that supplies the missing source trigger",
        ],
        "conclusion": conclusion,
        "records": records,
    }


def point_text(points: list[dict]) -> str:
    return ", ".join(f"{point['x']},{point['y']}" for point in points) or "-"


def html_page(summary: dict) -> str:
    reference_rows = []
    for row in summary["records"]:
        maps = ", ".join(row["nearbyMapResources"]) or "-"
        scene = row.get("sceneIdHex") if row.get("sceneRecordLike") else "-"
        reference_rows.append(
            "<tr>"
            f"<td><code>{html.escape(row['filename'])}</code></td>"
            f"<td><code>{html.escape(row['refVaHex'])}</code></td>"
            f"<td><code>{html.escape(scene)}</code></td>"
            f"<td>{html.escape(maps)}</td>"
            f"<td>{row['pointCandidateCount']}</td>"
            f"<td>{html.escape(str(row.get('pointCandidateClassCounts') or {}))}</td>"
            f"<td>{row['strictSourceTargetCandidateCount']}</td>"
            "</tr>"
        )
    point_rows = []
    for row in summary["records"]:
        for point in row["pointCandidates"]:
            owner = point.get("ownerSceneRecord") or {}
            owner_label = f"{owner.get('mapStem')} {owner.get('recordVaHex')}" if owner else "-"
            point_rows.append(
                "<tr>"
                f"<td><code>{html.escape(row['filename'])}</code><br><code>{html.escape(row['refVaHex'])}</code></td>"
                f"<td><code>{html.escape(point['fieldVaHex'])}</code></td>"
                f"<td><code>{html.escape(point['payloadVaHex'])}</code></td>"
                f"<td>{html.escape(str(point.get('classification')))}</td>"
                f"<td>{html.escape(owner_label)}</td>"
                f"<td>{point['rawPointCount']}</td>"
                f"<td>{point['inBoundsPointCount']}</td>"
                f"<td>{point['exactRouteExitHitCount']}</td>"
                f"<td>{html.escape(point_text(point['firstRawPoints']))}</td>"
                f"<td>{point['pointerRefCount']}</td>"
                "</tr>"
            )
    proof_items = "".join(f"<li>{html.escape(item)}</li>" for item in summary["remainingProofs"])
    return "\n".join([
        "<!doctype html>",
        '<html lang="en">',
        "<head>",
        '  <meta charset="utf-8">',
        '  <meta name="viewport" content="width=device-width, initial-scale=1">',
        "  <title>map1_01a Resource Reference Scan</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>map1_01a Resource Reference Scan</h1>",
        f"  <p>Route: <code>{html.escape(summary['source'])} -&gt; {html.escape(summary['target'])}</code>. "
        f"Resource refs: {summary['resourceReferenceCount']}; point candidates: {summary['pointCandidateCount']}; "
        f"point candidate classes: {html.escape(str(summary.get('pointCandidateClassCounts') or {}))}; "
        f"route-exit point candidates: {summary['routeExitPointCandidateCount']}; "
        f"strict source-target candidates: {summary['strictSourceTargetCandidateCount']}; "
        "current frontier classes: "
        f"{html.escape(str(summary.get('currentFrontierPointCandidateClassCounts') or {}))}; "
        f"current frontier route hits: {summary.get('currentFrontierRouteExitPointHitCount')}.</p>",
        f"  <p>{html.escape(summary['conclusion'])}</p>",
        "  <h2>Resource References</h2>",
        "  <table><thead><tr><th>resource</th><th>ref</th><th>scene</th><th>nearby maps</th><th>point candidates</th><th>classes</th><th>strict</th></tr></thead>",
        f"  <tbody>{''.join(reference_rows) or '<tr><td colspan=\"7\">No refs.</td></tr>'}</tbody></table>",
        "  <h2>Point Candidates</h2>",
        "  <table><thead><tr><th>resource ref</th><th>field</th><th>payload</th><th>class</th><th>owner</th><th>raw</th><th>in bounds</th><th>route hits</th><th>first points</th><th>refs</th></tr></thead>",
        f"  <tbody>{''.join(point_rows) or '<tr><td colspan=\"10\">No point candidates.</td></tr>'}</tbody></table>",
        "  <h2>Remaining Proofs</h2>",
        f"  <ul>{proof_items}</ul>",
        "</body>",
        "</html>",
        "",
    ])


def write_outputs(summary: dict, out_dir: Path = OUT) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "map1_01a_resource_ref_scan.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("--maps", type=Path, default=OUT / "maps.js")
    parser.add_argument("--map-exit-candidates", type=Path, default=OUT / "map_exit_candidates.json")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    summary = build_summary(
        args.exe.read_bytes(),
        parse_maps_js(args.maps),
        load_json(args.map_exit_candidates, []),
    )
    write_outputs(summary, args.out_dir)
    print(f"wrote map1_01a resource reference scan -> {args.out_dir / 'map1_01a_resource_ref_scan.json'}")


if __name__ == "__main__":
    main()