#!/usr/bin/env python3
"""Export decompressed CNS map layouts as a browser-loadable JS asset."""
from __future__ import annotations

import argparse
import json
import struct
from pathlib import Path

from decode_cns import decompress_cns, indices_to_rgb, parse_image
from probe_exe_scene_tables import read_sections, va_to_offset


ROOT = Path(__file__).resolve().parents[1]
EXE = ROOT / "Hwanse2.exe"
LOCATION_SELECTOR_TABLE_VA = 0x004EBC14

LAYER_TILESET_OVERRIDES: dict[str, list[str]] = {}

SCENE_TILESET_OVERRIDES = {
    "map1_01a": ["map_b1", "map_b2", "map_b3"],
    "map1_02b": ["map_f1", "map_f2", "map_f3"],
    "map3_09a": ["map_d1", "map_d2", "map_d3"],
    "map5_07h": ["map_h1", "map_h2", "map_h3"],
    "map5_38i": ["map_i1", "map_i2", "map_i3"],
    "map2_03l": ["map_l1", "map_l2"],
    "map4_08n": ["map_n1", "map_n2", "map_n3"],
    "map7_02f": ["map_h1", "map_h2", "map_h3"],
    "map7_03h": ["map_h1", "map_h2", "map_h3"],
    "map8_18o": ["map_o1", "map_o2", "map_o3"],
    "map8_32q": ["map_q1", "map_q2", "map_q3"],
    "map9_01e": ["map_e1", "map_e2", "map_e3"],
}

SCENE_ID_OVERRIDES = {
    "map1_01a": 0x0518,
    "map1_02b": 0x0618,
    "map3_09a": 0x3518,
    "map5_07h": 0x4618,
    "map5_38i": 0x6518,
    "map2_03l": 0x0918,
    "map4_08n": 0x3E18,
    "map7_02f": 0x8318,
    "map7_03h": 0x8418,
    "map8_18o": 0x9E18,
    "map8_32q": 0xAC18,
    "map9_01e": 0xAD18,
}

PASSABLE_TILE_OVERRIDES: dict[str, list[int]] = {}

TRANSITION_OVERRIDES: dict[str, list[dict]] = {}

TILESET_LAYOUT_OVERRIDES: dict[str, dict] = {}

TERRAIN_BLOCKED_TILE_CACHE: dict[tuple[Path, str], list[int]] = {}
LOCATION_LABEL_CACHE: dict[int, str] | None = None

EMPTY_LAYER1_TILES = [1, 2]


def cp949_text_atom(exe: bytes, sections: list[dict], va: int, limit: int = 96) -> str:
    offset = va_to_offset(sections, va)
    if offset is None:
        return ""
    raw = bytearray()
    for index in range(offset, min(len(exe), offset + limit)):
        byte = exe[index]
        if byte == 0:
            break
        if byte == 0x40 and index + 1 < len(exe) and 0 <= exe[index + 1] < 0x40:
            break
        raw.append(byte)
    return bytes(raw).decode("cp949", errors="replace")


def location_labels() -> dict[int, str]:
    global LOCATION_LABEL_CACHE
    if LOCATION_LABEL_CACHE is not None:
        return LOCATION_LABEL_CACHE
    labels: dict[int, str] = {}
    if not EXE.exists():
        LOCATION_LABEL_CACHE = labels
        return labels
    exe = EXE.read_bytes()
    sections = read_sections(exe)
    offset = va_to_offset(sections, LOCATION_SELECTOR_TABLE_VA)
    if offset is None:
        LOCATION_LABEL_CACHE = labels
        return labels
    count = struct.unpack_from("<I", exe, offset + 4)[0]
    for index in range(count):
        pointer = struct.unpack_from("<I", exe, offset + 8 + index * 4)[0]
        labels[index] = cp949_text_atom(exe, sections, pointer)
    LOCATION_LABEL_CACHE = labels
    return labels


def terrain_blocked_tiles(src_dir: Path, tileset_name: str) -> list[int]:
    key = (src_dir, tileset_name)
    if key in TERRAIN_BLOCKED_TILE_CACHE:
        return TERRAIN_BLOCKED_TILE_CACHE[key]

    src = src_dir / f"{tileset_name}.cns"
    decoded = decompress_cns(src.read_bytes())
    width, height, palette, pixels, bpp = parse_image(decoded)
    rgb = indices_to_rgb(width, height, palette, pixels, bpp)
    tile_size = 16
    columns = width // tile_size
    rows = height // tile_size
    transparent = rgb[:3]
    blocked = []

    for tile_id in range(columns * rows):
        index = tile_id
        src_x = (index % columns) * tile_size
        src_y = (index // columns) * tile_size
        visible = []
        for y in range(tile_size):
            for x in range(tile_size):
                offset = ((src_y + y) * width + src_x + x) * 3
                color = rgb[offset : offset + 3]
                if color != transparent:
                    visible.append(color)
        if not visible:
            blocked.append(tile_id)
            continue
        brightness = sum(sum(color) / 3 for color in visible) / len(visible)
        if brightness < 45:
            blocked.append(tile_id)

    TERRAIN_BLOCKED_TILE_CACHE[key] = blocked
    return blocked

def parse_map(src: Path, name: str | None = None, metadata: dict | None = None) -> dict:
    decoded = decompress_cns(src.read_bytes())
    if len(decoded) < 4:
        raise ValueError(f"{src} is too short")

    data_offset = 4
    width = int.from_bytes(decoded[0:2], "little")
    height = int.from_bytes(decoded[2:4], "little")
    tile_count = width * height
    expected = 4 + tile_count * 4
    if len(decoded) != expected and len(decoded) >= 2:
        compact_width = decoded[0]
        compact_height = decoded[1]
        compact_tile_count = compact_width * compact_height
        compact_expected = 2 + compact_tile_count * 4
        if compact_width and compact_height and len(decoded) == compact_expected:
            data_offset = 2
            width = compact_width
            height = compact_height
            tile_count = compact_tile_count
            expected = compact_expected

    if len(decoded) != expected:
        raise ValueError(
            f"{src} decoded size {len(decoded)} does not match "
            f"{width}x{height} two-layer map size {expected}"
        )

    cell_values = [
        int.from_bytes(decoded[i : i + 2], "little")
        for i in range(data_offset, len(decoded), 2)
    ]
    # The body is stored as two planar u16 blocks. Treating adjacent values as
    # per-cell pairs splits one continuous map into repeated left/right chunks.
    layers = [
        cell_values[:tile_count],
        cell_values[tile_count:],
    ]

    map_name = name or src.stem
    suffix = map_name[-1] if map_name[-1:].isalpha() else ""
    default_tileset = f"map_{suffix}1" if suffix else "map_a1"
    default_second_tileset = f"map_{suffix}2" if suffix else default_tileset
    metadata = metadata or {}
    metadata_tilesets = metadata.get("tilesets") or []
    layer_tilesets = LAYER_TILESET_OVERRIDES.get(
        map_name,
        [default_tileset, default_second_tileset],
    )
    observed_scene_tilesets = metadata_tilesets or SCENE_TILESET_OVERRIDES.get(map_name, layer_tilesets)
    variant_candidates = [
        layer_tilesets,
        *(metadata.get("tilesetVariants") or []),
    ]
    scene_tileset_variants = []
    for variant in variant_candidates:
        key = tuple(variant)
        if key and key not in {tuple(item) for item in scene_tileset_variants}:
            scene_tileset_variants.append(list(variant))
    scene_id = metadata.get("sceneId", SCENE_ID_OVERRIDES.get(map_name))
    location_index = scene_id >> 8 if isinstance(scene_id, int) else None
    label_table = location_labels()
    location_label = label_table.get(location_index, "") if location_index is not None else ""
    result = {
        "name": map_name,
        "tileset": layer_tilesets[0],
        "layerTilesets": layer_tilesets,
        "sceneTilesets": layer_tilesets,
        "observedSceneTilesets": observed_scene_tilesets,
        "sceneId": scene_id,
        "sceneIdHex": metadata.get("sceneIdHex"),
        "sceneRecordVa": metadata.get("recordVa"),
        "locationIndex": location_index,
        "locationLabel": location_label,
        "locationLabelSource": "sceneIdHighByte->0x004ebc14" if location_label else "",
        "sceneVariantCount": metadata.get("variantCount"),
        "sceneTilesetVariants": scene_tileset_variants,
        "passableTiles": PASSABLE_TILE_OVERRIDES.get(map_name),
        "terrainBlockedTiles": terrain_blocked_tiles(src.parent, layer_tilesets[0]),
        "emptyLayer1Tiles": EMPTY_LAYER1_TILES,
        "transitions": TRANSITION_OVERRIDES.get(map_name, []),
        "width": width,
        "height": height,
        "tileSize": 16,
        "tileIndexOffset": 0,
        "layers": layers,
    }
    result.update(
        TILESET_LAYOUT_OVERRIDES.get(
            map_name,
            {
                "tilesetColumns": 40,
            },
        )
    )
    return result


def export_maps(srcs: list[Path], dst: Path, metadata_by_map: dict[str, dict] | None = None) -> None:
    metadata_by_map = metadata_by_map or {}
    maps = {
        src.stem: parse_map(src, metadata=metadata_by_map.get(src.stem))
        for src in srcs
    }
    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.write_text(
        "window.HWANSE_MAPS = "
        + json.dumps(maps, ensure_ascii=False, separators=(",", ":"))
        + ";\n",
        encoding="utf-8",
    )
    print(f"wrote {len(maps)} maps -> {dst}")
    for name, info in maps.items():
        print(f"  {name}: {info['width']}x{info['height']}, tileset={info['tileset']}")


def map_index_entry(info: dict) -> dict:
    return {
        key: value
        for key, value in info.items()
        if key not in {"layers", "terrainBlockedTiles", "passableTiles", "transitions"}
    } | {"layerCount": len(info.get("layers") or [])}


def export_map_chunks(srcs: list[Path], index_dst: Path, chunk_dir: Path, metadata_by_map: dict[str, dict] | None = None) -> None:
    metadata_by_map = metadata_by_map or {}
    maps = {
        src.stem: parse_map(src, metadata=metadata_by_map.get(src.stem))
        for src in srcs
    }
    index_dst.parent.mkdir(parents=True, exist_ok=True)
    chunk_dir.mkdir(parents=True, exist_ok=True)
    for old in chunk_dir.glob("*.js"):
        old.unlink()

    index = {name: map_index_entry(info) for name, info in maps.items()}
    index_dst.write_text(
        "window.HWANSE_MAPS = window.HWANSE_MAPS || {};"
        "window.HWANSE_MAP_INDEX = "
        + json.dumps(index, ensure_ascii=False, separators=(",", ":"))
        + ";\n",
        encoding="utf-8",
    )
    for name, info in maps.items():
        (chunk_dir / f"{name}.js").write_text(
            "window.HWANSE_MAPS = window.HWANSE_MAPS || {};"
            f"window.HWANSE_MAPS[{json.dumps(name)}] = "
            + json.dumps(info, ensure_ascii=False, separators=(",", ":"))
            + ";\n",
            encoding="utf-8",
        )
    print(f"wrote {len(maps)} runtime map chunks -> {chunk_dir}")
    print(f"wrote runtime map index -> {index_dst}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("src", nargs="+", type=Path)
    parser.add_argument("dst", type=Path)
    args = parser.parse_args()

    export_maps(args.src, args.dst)


if __name__ == "__main__":
    main()