#!/usr/bin/env python3
"""Dump probable CNS scene-load records embedded in Hwanse2.exe.

This is a diagnostic helper for the map renderer. The executable contains
arrays of small records that reference CNS filenames plus extra arguments.
Those extra arguments are likely needed to pick tilesets, source tables, map
logic, or scene scripts; treating map CNS values as direct tile IDs has proven
insufficient.
"""
from __future__ import annotations

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


IMAGE_BASE = 0x400000


def read_sections(exe: bytes) -> list[dict]:
    if exe[:2] != b"MZ":
        raise ValueError("not an MZ/PE executable")
    pe_off = struct.unpack_from("<I", exe, 0x3C)[0]
    if exe[pe_off : pe_off + 4] != b"PE\0\0":
        raise ValueError("not a PE executable")
    section_count = struct.unpack_from("<H", exe, pe_off + 6)[0]
    optional_size = struct.unpack_from("<H", exe, pe_off + 20)[0]
    section_off = pe_off + 24 + optional_size
    sections = []
    for index in range(section_count):
        off = section_off + index * 40
        name = exe[off : off + 8].split(b"\0", 1)[0].decode("ascii", "replace")
        virtual_size, virtual_address, raw_size, raw_pointer = struct.unpack_from(
            "<IIII", exe, off + 8
        )
        sections.append(
            {
                "name": name,
                "va": IMAGE_BASE + virtual_address,
                "size": max(virtual_size, raw_size),
                "raw_size": raw_size,
                "raw": raw_pointer,
            }
        )
    return sections


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


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


def c_string(exe: bytes, offset: int) -> str:
    end = exe.find(b"\0", offset)
    if end < 0:
        end = len(exe)
    return exe[offset:end].decode("ascii", "replace")


def find_cns_strings(exe: bytes, sections: list[dict]) -> dict[int, str]:
    strings: dict[int, str] = {}
    for match in re.finditer(rb"[a-z0-9_]{2,12}\.cns\0", exe):
        va = offset_to_va(sections, match.start())
        if va is not None:
            strings[va] = c_string(exe, match.start())
    return strings


def format_dword(value: int, strings: dict[int, str]) -> str:
    if value in strings:
        return f"{value:08x} -> {strings[value]}"
    if 0x00400000 <= value <= 0x00600000:
        return f"{value:08x} -> ptr?"
    return f"{value:08x}"


def classify_cns(name: str) -> str:
    if re.fullmatch(r"map_[a-z][123]\.cns", name):
        return "tileset"
    if re.fullmatch(r"map\d+_\d+[a-z]\.cns", name):
        return "map"
    if name.startswith("cara_") or name.startswith("face_"):
        return "sprite"
    return "other"


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


def collect_reference(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    hit: int,
    context: int,
) -> dict | None:
    hit_va = offset_to_va(sections, hit)
    if hit_va is None:
        return None

    values = []
    aligned_base = max(0, hit - context * 4)
    aligned_base -= aligned_base % 4
    end = min(len(exe), hit + (context + 1) * 4)
    for off in range(aligned_base, end, 4):
        if off + 4 > len(exe):
            break
        va = offset_to_va(sections, off)
        if va is None:
            continue
        value = dword_at(exe, off)
        item = {"va": va, "value": value, "isTarget": off == hit}
        if value in strings:
            item["cns"] = strings[value]
            item["kind"] = classify_cns(strings[value])
        elif 0x00400000 <= value <= 0x00600000:
            item["pointer"] = True
        values.append(item)

    after = [item for item in values if item["va"] > hit_va and "cns" in item]
    return {
        "fileOffset": hit,
        "va": hit_va,
        "values": values,
        "nearbyCns": [{"name": item["cns"], "kind": item["kind"]} for item in after],
        "nearbyTilesets": [
            item["cns"][:-4]
            for item in after
            if item.get("kind") == "tileset"
        ],
    }


def collect_records(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    target: str,
    context: int,
) -> list[dict]:
    target_vas = [va for va, name in strings.items() if target.lower() in name.lower()]
    if not target_vas:
        raise ValueError(f"no CNS string matching {target!r}")

    records = []
    for target_va in target_vas:
        pattern = struct.pack("<I", target_va)
        starts = []
        search = 0
        while True:
            hit = exe.find(pattern, search)
            if hit < 0:
                break
            starts.append(hit)
            search = hit + 1

        refs = []
        for hit in starts:
            ref = collect_reference(exe, sections, strings, hit, context)
            if ref is not None:
                refs.append(ref)
        records.append(
            {
                "name": strings[target_va],
                "va": target_va,
                "references": refs,
            }
        )
    return records


def dump_records(
    exe: bytes,
    sections: list[dict],
    strings: dict[int, str],
    target: str,
    context: int,
) -> None:
    records = collect_records(exe, sections, strings, target, context)
    for record in records:
        print(
            f"{record['name']} @ va=0x{record['va']:08x}: "
            f"{len(record['references'])} pointer reference(s)"
        )
        for ref in record["references"]:
            print(f"  ref file=0x{ref['fileOffset']:06x} va=0x{ref['va']:08x}")
            for item in ref["values"]:
                mark = "*" if item["isTarget"] else " "
                print(f"   {mark}0x{item['va']:08x}: {format_dword(item['value'], strings)}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--exe", type=Path, default=Path("Hwanse2.exe"))
    parser.add_argument("--contains", required=True, help="CNS filename substring")
    parser.add_argument("--context", type=int, default=8, help="dwords before/after each reference")
    parser.add_argument("--json", action="store_true", help="write structured records as JSON")
    args = parser.parse_args()

    exe = args.exe.read_bytes()
    sections = read_sections(exe)
    strings = find_cns_strings(exe, sections)
    if args.json:
        print(
            json.dumps(
                collect_records(exe, sections, strings, args.contains, args.context),
                ensure_ascii=False,
                indent=2,
            )
        )
    else:
        dump_records(exe, sections, strings, args.contains, args.context)


if __name__ == "__main__":
    main()