#!/usr/bin/env python3
"""Extract probable scene hotspot coordinate tables from Hwanse2.exe."""
from __future__ import annotations

import argparse
import json
import re
import struct
import sys
from pathlib import Path

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

from decode_cns import decompress_cns
from probe_exe_scene_tables import c_string, find_cns_strings, offset_to_va, read_sections, va_to_offset


MAP_RE = re.compile(r"map\d+_\d+[a-z]\.cns$")


def dword_at(data: bytes, offset: int) -> int:
    return struct.unpack_from("<I", data, offset)[0]


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


def find_pointer_refs(
    data: bytes,
    sections: list[dict],
    value: int,
    section_names: set[str] | None = None,
) -> list[dict]:
    needle = struct.pack("<I", value)
    refs = []
    search = 0
    while True:
        hit = data.find(needle, search)
        if hit < 0:
            break
        search = hit + 1
        section = section_for_offset(sections, hit)
        if section is None:
            continue
        if section_names is not None and section["name"] not in section_names:
            continue
        ref_va = offset_to_va(sections, hit)
        if ref_va is None:
            continue
        item = {
            "section": section["name"],
            "refVa": ref_va,
            "refFileOffset": hit,
            "targetVa": value,
        }
        if hit >= 4:
            previous_value = dword_at(data, hit - 4)
            item["previousValue"] = previous_value
            item["previousValueHex"] = f"0x{previous_value:08x}"
            if 0x00400000 <= previous_value <= 0x00600000:
                item["conditionVa"] = previous_value
                condition_dwords = read_tail_dwords(data, sections, previous_value, count=8)
                item["conditionFirstDwords"] = condition_dwords
                if len(condition_dwords) > 1 and condition_dwords[1].get("pointer"):
                    payload_va = condition_dwords[1]["value"]
                    payload_dwords = read_tail_dwords(data, sections, payload_va, count=24)
                    item["conditionPayloadVa"] = payload_va
                    item["conditionPayloadDwords"] = payload_dwords
                    item["conditionLinkedStrings"] = linked_strings_from_dwords(payload_dwords)
        refs.append(item)
    return refs


def load_map_sizes(extract_dir: Path) -> dict[str, tuple[int, int]]:
    sizes = {}
    for path in extract_dir.glob("map*.cns"):
        if not re.fullmatch(r"map\d+_\d+[a-z]\.cns", path.name):
            continue
        decoded = decompress_cns(path.read_bytes())
        if len(decoded) < 4:
            continue
        width = int.from_bytes(decoded[0:2], "little")
        height = int.from_bytes(decoded[2:4], "little")
        if len(decoded) != 4 + width * height * 4 and len(decoded) >= 2:
            compact_width = decoded[0]
            compact_height = decoded[1]
            if len(decoded) == 2 + compact_width * compact_height * 4:
                width = compact_width
                height = compact_height
        sizes[path.stem] = (width, height)
    return sizes


def read_point_table(data: bytes, sections: list[dict], ptr_va: int, width: int, height: int) -> tuple[list[dict], list[dict]]:
    offset = va_to_offset(sections, ptr_va)
    if offset is None:
        return [], []

    raw_points = []
    in_bounds = []
    for index in range(512):
        off = offset + index * 4
        if off + 4 > len(data):
            break
        value = dword_at(data, off)
        if value == 0:
            break
        x = value & 0xFFFF
        y = value >> 16
        if x >= 512 or y >= 512:
            break
        point = {"x": x, "y": y}
        raw_points.append(point)
        if x < width and y < height:
            in_bounds.append(point)
    return raw_points, in_bounds


def active_points_from_hint(raw_points: list[dict], width: int, height: int, point_count_hint: int) -> list[dict]:
    if point_count_hint <= 0 or point_count_hint > len(raw_points):
        candidates = raw_points
    else:
        candidates = raw_points[:point_count_hint]
    return [
        point
        for point in candidates
        if point["x"] < width and point["y"] < height
    ]


def printable_prefix(text: str) -> str | None:
    if not text:
        return None
    prefix = text[:32]
    if all(32 <= ord(char) < 127 for char in prefix):
        return prefix
    return None


def read_tail_dwords(data: bytes, sections: list[dict], start_va: int, count: int = 12) -> list[dict]:
    offset = va_to_offset(sections, start_va)
    if offset is None:
        return []

    values = []
    for index in range(count):
        off = offset + index * 4
        if off + 4 > len(data):
            break
        value = dword_at(data, off)
        item = {
            "va": start_va + index * 4,
            "value": value,
            "hex": f"0x{value:08x}",
            "lo": value & 0xFFFF,
            "hi": value >> 16,
        }
        if 0x00400000 <= value <= 0x00600000:
            item["pointer"] = True
            ptr_offset = va_to_offset(sections, value)
            if ptr_offset is not None:
                text = printable_prefix(c_string(data, ptr_offset))
                if text:
                    item["string"] = text
        values.append(item)
    return values


def read_pointer_blocks(data: bytes, sections: list[dict], tail_dwords: list[dict]) -> list[dict]:
    blocks = []
    for item in tail_dwords:
        if not item.get("pointer"):
            continue
        ptr_va = item["value"]
        values = read_tail_dwords(data, sections, ptr_va, count=16)
        first_value = values[0]["value"] if values else None
        pointer_count = sum(1 for value in values if value.get("pointer"))
        strings = [value["string"] for value in values if value.get("string")]
        if first_value == 0x3F:
            block_kind = "selector"
        elif first_value == 0x42:
            block_kind = "pointerTable"
        elif first_value == 0x08:
            block_kind = "scriptHeader"
        elif first_value == 0x01:
            block_kind = "scriptEntry"
        elif first_value == 0x19:
            block_kind = "prefixedSelector"
        elif first_value in {0x03, 0x16D}:
            block_kind = "scriptLike"
        else:
            block_kind = "unknown"
        blocks.append(
            {
                "pointerVa": ptr_va,
                "sourceVa": item["va"],
                "blockKind": block_kind,
                "firstValue": first_value,
                "pointerCount": pointer_count,
                "strings": strings,
                "firstDwords": values,
            }
        )
    return blocks


def collect_linked_strings(blocks: list[dict]) -> list[str]:
    seen = set()
    linked = []
    for block in blocks:
        for text in block["strings"]:
            if text in seen:
                continue
            seen.add(text)
            linked.append(text)
    return linked


def linked_strings_from_dwords(values: list[dict]) -> list[str]:
    seen = set()
    linked = []
    for item in values:
        text = item.get("string")
        if not text or text in seen:
            continue
        seen.add(text)
        linked.append(text)
    return linked


def extract_events(data: bytes, extract_dir: Path) -> list[dict]:
    sections = read_sections(data)
    strings = find_cns_strings(data, sections)
    map_sizes = load_map_sizes(extract_dir)
    records = []

    for map_va, filename in strings.items():
        if not MAP_RE.fullmatch(filename):
            continue
        name = filename[:-4]
        size = map_sizes.get(name)
        if not size:
            continue

        search = 0
        needle = struct.pack("<I", map_va)
        while True:
            hit = data.find(needle, search)
            if hit < 0:
                break
            search = hit + 1
            if hit + 28 > len(data):
                continue

            scene_id = dword_at(data, hit + 4)
            zero = dword_at(data, hit + 8)
            event_kind = dword_at(data, hit + 12)
            sentinel = dword_at(data, hit + 16)
            point_ptr = dword_at(data, hit + 20)
            point_count_hint = dword_at(data, hit + 24)
            if zero != 0 or sentinel != 0x3F:
                continue
            if not (0x00400000 <= point_ptr <= 0x00600000):
                continue

            width, height = size
            raw_points, in_bounds_points = read_point_table(data, sections, point_ptr, width, height)
            active_points = active_points_from_hint(raw_points, width, height, point_count_hint)
            if not in_bounds_points:
                continue
            record_va = offset_to_va(sections, hit)
            if record_va is None:
                continue
            point_table_end_va = point_ptr + len(raw_points) * 4
            tail_dwords = read_tail_dwords(data, sections, point_table_end_va)
            tail_pointer_blocks = read_pointer_blocks(data, sections, tail_dwords)
            dispatch_refs = find_pointer_refs(data, sections, record_va + 12, {".data", ".rdata"})
            records.append(
                {
                    "map": name,
                    "sceneId": scene_id,
                    "sceneIdHex": f"0x{scene_id:04x}",
                    "recordVa": record_va,
                    "recordFileOffset": hit,
                    "eventKind": event_kind,
                    "eventDispatchVa": record_va + 12,
                    "eventDispatchRefs": dispatch_refs,
                    "pointTableVa": point_ptr,
                    "pointTableEndVa": point_table_end_va,
                    "pointCountHint": point_count_hint,
                    "rawPoints": raw_points,
                    "rawPointCount": len(raw_points),
                    "inBoundsPoints": in_bounds_points,
                    "inBoundsPointCount": len(in_bounds_points),
                    "activePointCount": len(active_points),
                    "allRawPointsInBounds": len(raw_points) == len(in_bounds_points),
                    "points": in_bounds_points,
                    "activePoints": active_points,
                    "linkedStrings": collect_linked_strings(tail_pointer_blocks),
                    "tailDwords": tail_dwords,
                    "tailPointerBlocks": tail_pointer_blocks,
                }
            )

    records.sort(key=lambda item: (item["recordVa"], item["map"]))
    return records


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--exe", type=Path, default=Path("Hwanse2.exe"))
    parser.add_argument("--extract-dir", type=Path, default=Path("extract_fld"))
    parser.add_argument("--out", type=Path, default=Path("out/scene_events.json"))
    parser.add_argument("--map", dest="map_filter", help="only print records for one map stem")
    args = parser.parse_args()

    events = extract_events(args.exe.read_bytes(), args.extract_dir)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(events, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    shown = [item for item in events if not args.map_filter or item["map"] == args.map_filter]
    print(f"wrote {len(events)} scene event records -> {args.out}")
    for item in shown[:100]:
        print(
            f"{item['sceneIdHex']} {item['map']} "
            f"kind={item['eventKind']} active={len(item.get('activePoints', item['points']))} "
            f"points={len(item['points'])} "
            f"record=0x{item['recordVa']:08x}"
        )
    if len(shown) > 100:
        print(f"... {len(shown) - 100} more")


if __name__ == "__main__":
    main()