#!/usr/bin/env python3
"""Slice UI Normal atlas into Common single sprites."""
from __future__ import annotations

import json
from collections import deque
from pathlib import Path

import numpy as np
from PIL import Image

ATLAS = Path(r"d:\UD\AI\AIGame\Tools\output\ui_normal_atlas_gpt.png")
OUT_DIR = Path(r"d:\UD\AI\AIGame\AIClient\Assets\Bundles\UI\Texture\Common")
REPORT = Path(r"d:\UD\AI\AIGame\Tools\output\ui_normal_slice_report.json")

# From Doc/UI_Normal_Set.md (left-top origin)
SLICES = [
    ("ui_panel_main", 8, 8, 480, 256),
    ("ui_panel_popup", 496, 8, 360, 240),
    ("ui_panel_content", 864, 8, 320, 200),
    ("ui_panel_header", 1192, 8, 400, 80),
    ("ui_overlay_dim", 1600, 8, 192, 192),
    ("ui_divider_h", 1192, 96, 400, 16),
    ("btn_primary_normal", 8, 280, 224, 88),
    ("btn_primary_pressed", 240, 280, 224, 88),
    ("btn_primary_disabled", 472, 280, 224, 88),
    ("btn_secondary_normal", 704, 280, 224, 88),
    ("btn_secondary_pressed", 936, 280, 224, 88),
    ("btn_small_normal", 1168, 280, 160, 64),
    ("btn_small_pressed", 1336, 280, 160, 64),
    ("btn_close", 1504, 280, 72, 72),
    ("btn_back", 1584, 280, 72, 72),
    ("scrollbar_track_v", 1664, 280, 28, 160),
    ("scrollbar_thumb_v", 1700, 280, 28, 80),
    ("input_bg_normal", 8, 400, 360, 72),
    ("input_bg_focus", 376, 400, 360, 72),
    ("ui_list_item_normal", 744, 400, 320, 72),
    ("ui_list_item_selected", 1072, 400, 320, 72),
    ("tab_active", 1400, 400, 200, 72),
    ("tab_inactive", 1608, 400, 200, 72),
    ("progress_bg", 8, 504, 400, 32),
    ("progress_fill", 416, 504, 400, 32),
    ("slider_track", 824, 504, 320, 28),
    ("slider_fill", 1152, 504, 320, 28),
    ("slider_handle", 1480, 496, 56, 56),
    ("toast_bg", 1544, 496, 280, 80),
    ("toggle_bg_off", 8, 616, 112, 56),
    ("toggle_bg_on", 128, 616, 112, 56),
    ("toggle_knob", 248, 616, 48, 48),
    ("checkbox_off", 304, 616, 56, 56),
    ("checkbox_on", 368, 616, 56, 56),
    ("ico_search", 432, 616, 56, 56),
    ("ico_check", 496, 616, 56, 56),
    ("ico_settings", 560, 616, 56, 56),
    ("ico_refresh", 624, 616, 56, 56),
]


def find_boxes(arr: np.ndarray, alpha_thr: int = 24, min_px: int = 60):
    mask = arr[:, :, 3] > alpha_thr
    h, w = mask.shape
    visited = np.zeros_like(mask, dtype=bool)
    boxes = []
    for y in range(h):
        for x in np.where(mask[y] & ~visited[y])[0]:
            if visited[y, x] or not mask[y, x]:
                continue
            q = deque([(int(y), int(x))])
            visited[y, x] = True
            minx = maxx = int(x)
            miny = maxy = int(y)
            cnt = 0
            while q:
                cy, cx = q.popleft()
                cnt += 1
                for ny, nx in ((cy - 1, cx), (cy + 1, cx), (cy, cx - 1), (cy, cx + 1)):
                    if 0 <= ny < h and 0 <= nx < w and mask[ny, nx] and not visited[ny, nx]:
                        visited[ny, nx] = True
                        q.append((ny, nx))
                        minx = min(minx, nx)
                        maxx = max(maxx, nx)
                        miny = min(miny, ny)
                        maxy = max(maxy, ny)
            bw, bh = maxx - minx + 1, maxy - miny + 1
            if cnt >= min_px and bw >= 10 and bh >= 8:
                boxes.append(
                    {
                        "miny": miny,
                        "minx": minx,
                        "maxy": maxy,
                        "maxx": maxx,
                        "w": bw,
                        "h": bh,
                        "cx": (minx + maxx) / 2,
                        "cy": (miny + maxy) / 2,
                        "px": cnt,
                    }
                )
    return boxes


def match_slices(boxes: list[dict]) -> list[dict]:
    used = set()
    matches = []
    for name, x, y, w, h in SLICES:
        pcx, pcy = x + w / 2, y + h / 2
        best_i, best_score = None, 1e18
        for i, b in enumerate(boxes):
            if i in used:
                continue
            # prefer position; lightly penalize size mismatch
            dist = (b["cx"] - pcx) ** 2 + (b["cy"] - pcy) ** 2
            size_pen = ((b["w"] - w) / max(w, 1)) ** 2 + ((b["h"] - h) / max(h, 1)) ** 2
            score = dist + size_pen * 8000
            if score < best_score:
                best_score = score
                best_i = i
        if best_i is None:
            raise RuntimeError(f"no box for {name}")
        used.add(best_i)
        b = boxes[best_i]
        matches.append(
            {
                "name": name,
                "planned": [x, y, w, h],
                "box": [b["minx"], b["miny"], b["w"], b["h"]],
                "score": round(best_score, 1),
            }
        )
    return matches


def main() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    img = Image.open(ATLAS).convert("RGBA")
    arr = np.array(img)
    boxes = find_boxes(arr)
    if len(boxes) != len(SLICES):
        raise SystemExit(f"expected {len(SLICES)} components, got {len(boxes)}")

    matches = match_slices(boxes)
    written = []
    for m in matches:
        x, y, w, h = m["box"]
        crop = img.crop((x, y, x + w, y + h))
        # pad 2px transparent to reduce bleed when bilinear sampling
        pad = 2
        canvas = Image.new("RGBA", (w + pad * 2, h + pad * 2), (0, 0, 0, 0))
        canvas.paste(crop, (pad, pad))
        out = OUT_DIR / f"{m['name']}.png"
        canvas.save(out)
        written.append({"name": m["name"], "path": str(out), **m})

    REPORT.write_text(json.dumps({"count": len(written), "items": written}, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"sliced {len(written)} sprites -> {OUT_DIR}")
    print(f"report -> {REPORT}")
    for m in matches[:5]:
        print(m["name"], "planned", m["planned"], "->", m["box"], "score", m["score"])
    print("...")
    for m in matches[-3:]:
        print(m["name"], "planned", m["planned"], "->", m["box"], "score", m["score"])


if __name__ == "__main__":
    main()