#!/usr/bin/env python3
"""
Stitch WoWMapUprez_Dungeons 4×3 BLP tile grids and emit BOTH a hi-res PNG
(lossless intermediate) and a web-ready WebP per dungeon-floor in parallel.

Inputs:  data/uprez/Interface/Worldmap/<MapName>/<MapName><floor>_<idx>.blp
Outputs: data/maps_png_hires/<MapName>[_floor<n>].png   (lossless, ~8 MB)
         web/assets/maps/<MapName>[_floor<n>].webp      (q=82, ~600 KB-1.5 MB)
"""
from __future__ import annotations
import argparse
import os
import re
import sys
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from PIL import Image

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "data" / "uprez" / "Interface" / "Worldmap"
DST_PNG = ROOT / "data" / "maps_png_hires"
DST_WEBP = ROOT / "web" / "assets" / "maps"

# Two tile-naming conventions:
#   <Map><floor>_<idx>.blp   — multi-level dungeons (e.g. Naxxramas2_5.blp)
#   <Map><idx>.blp           — single-level dungeons (e.g. ZulFarrak5.blp)
# In both cases, idx is 1-12 (12 tiles per level in a 4×3 grid).
TILE_IDX_RE = re.compile(r"^(?P<prefix>.+?)(?P<idx>1[0-2]|[1-9])\.blp$", re.IGNORECASE)
FLOOR_SUFFIX_RE = re.compile(r"^(?P<map>.+?)(?P<floor>\d+)_$")
COLS, ROWS = 4, 3
TILES_PER_FLOOR = COLS * ROWS


def parse_tile(filename: str):
    """Return (mapname, floor, idx) or None. floor=0 indicates the
    "single-level" form with no explicit floor in the filename."""
    m = TILE_IDX_RE.match(filename)
    if not m:
        return None
    prefix = m.group("prefix")
    idx = int(m.group("idx"))
    fm = FLOOR_SUFFIX_RE.match(prefix)
    if fm:
        return fm.group("map"), int(fm.group("floor")), idx
    return prefix, 0, idx


def stitch_one(args):
    out_basename, tile_paths, png_dir, webp_dir, max_width, quality, write_png = args
    # Open tile #1 to learn dimensions; assume uniform.
    with Image.open(tile_paths[1]) as t1:
        tw, th = t1.size
    canvas = Image.new("RGBA", (COLS * tw, ROWS * th))
    for idx in range(1, TILES_PER_FLOOR + 1):
        with Image.open(tile_paths[idx]) as tile:
            col = (idx - 1) % COLS
            row = (idx - 1) // COLS
            canvas.paste(tile.convert("RGBA"), (col * tw, row * th))

    if write_png:
        canvas.save(png_dir / f"{out_basename}.png", optimize=False)

    # WebP needs RGB and may want a downscale
    img = canvas.convert("RGB")
    if max_width and img.width > max_width:
        ratio = max_width / img.width
        img = img.resize((max_width, int(img.height * ratio)), Image.LANCZOS)
    img.save(webp_dir / f"{out_basename}.webp", "WEBP", quality=quality, method=4)
    return out_basename, canvas.size


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-width", type=int, default=4096,
                    help="downscale WebP to this max width (default: keep 4096)")
    ap.add_argument("--quality", type=int, default=82)
    ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) - 1))
    ap.add_argument("--no-png", action="store_true",
                    help="skip lossless PNG intermediate (save disk + time)")
    args = ap.parse_args()

    if not SRC.exists():
        print(f"missing {SRC} — clone WoWMapUprez_Dungeons first", file=sys.stderr)
        return 1
    DST_PNG.mkdir(parents=True, exist_ok=True)
    DST_WEBP.mkdir(parents=True, exist_ok=True)

    # Group: { (mapname, floor): {idx: path} }; floor=0 is "no floor in filename"
    groups: dict[tuple[str, int], dict[int, Path]] = defaultdict(dict)
    for dungeon_dir in sorted(SRC.iterdir()):
        if not dungeon_dir.is_dir():
            continue
        for blp in dungeon_dir.glob("*.blp"):
            parsed = parse_tile(blp.name)
            if not parsed:
                continue
            mapname, floor, idx = parsed
            groups[(mapname, floor)][idx] = blp

    floors_per_map: dict[str, set[int]] = defaultdict(set)
    for (mapname, floor) in groups:
        floors_per_map[mapname].add(floor)

    jobs = []
    skipped = []
    for (mapname, floor), tiles in sorted(groups.items()):
        if len(tiles) != TILES_PER_FLOOR:
            skipped.append(f"{mapname} floor {floor}: {len(tiles)}/{TILES_PER_FLOOR}")
            continue
        floors = floors_per_map[mapname]
        if len(floors) == 1:
            suffix = ""                                       # only one variant
        elif floor == 0:
            suffix = "_overview"                              # mixed: 0 + others
        else:
            suffix = f"_floor{floor}"
        out_base = f"{mapname}{suffix}"
        jobs.append((out_base, tiles, DST_PNG, DST_WEBP, args.max_width, args.quality, not args.no_png))

    print(f"stitching {len(jobs)} dungeon-floors with {args.workers} workers...")
    done = 0
    with ProcessPoolExecutor(max_workers=args.workers) as pool:
        futures = [pool.submit(stitch_one, j) for j in jobs]
        for fut in as_completed(futures):
            try:
                base, size = fut.result()
                done += 1
                if done % 20 == 0 or done == len(jobs):
                    print(f"  {done}/{len(jobs)}  (latest: {base} {size[0]}×{size[1]})")
            except Exception as e:
                print(f"  FAIL: {e}", file=sys.stderr)

    print(f"done: {done}/{len(jobs)} stitched")
    if skipped:
        print(f"\nskipped {len(skipped)} (incomplete tile sets):", file=sys.stderr)
        for s in skipped[:10]:
            print(f"  {s}", file=sys.stderr)
        if len(skipped) > 10:
            print(f"  ... and {len(skipped) - 10} more", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())