#!/usr/bin/env python3
"""Deterministically split a GPT generated white-background parts sheet.

The model call is intentionally kept in Genarrative's Node runner.  This file
only performs the open-source, deterministic connected-components step used by
the Genielabs Spine Animation AI workflow.
"""
import argparse
import json
import os
import sys
import numpy as np

try:
    import cv2
except ImportError as error:
    raise SystemExit("缺少 OpenCV，请先执行：python3 -m pip install opencv-python") from error


def segment(atlas_path: str, output_dir: str, min_area: int, padding: int, bg_threshold: int):
    image = cv2.imread(atlas_path, cv2.IMREAD_UNCHANGED)
    if image is None:
        raise ValueError(f"无法读取拆件图集：{atlas_path}")
    if image.ndim == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGRA)
    elif image.shape[2] == 3:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)

    # Removing every bright pixel also erased eye whites, teeth and pale
    # clothing. Treat white connected to the canvas boundary as background;
    # enclosed highlights remain part of their surrounding cutout.
    white = np.all(image[:, :, :3] >= bg_threshold, axis=2)
    transparent = image[:, :, 3] == 0
    bg_candidates = (white | transparent).astype(np.uint8)
    _, bg_labels, _, _ = cv2.connectedComponentsWithStats(bg_candidates, 8)
    border_labels = np.unique(np.concatenate((bg_labels[0], bg_labels[-1], bg_labels[:, 0], bg_labels[:, -1])))
    border_labels = border_labels[border_labels != 0]
    background = np.isin(bg_labels, border_labels) | transparent
    mask = (~background).astype(np.uint8) * 255
    count, labels, stats, _ = cv2.connectedComponentsWithStats(mask, 8)
    os.makedirs(output_dir, exist_ok=True)
    height, width = image.shape[:2]
    saved = []
    components = []
    for label in range(1, count):
        area = int(stats[label, cv2.CC_STAT_AREA])
        if area < min_area:
            continue
        x, y = int(stats[label, cv2.CC_STAT_LEFT]), int(stats[label, cv2.CC_STAT_TOP])
        w, h = int(stats[label, cv2.CC_STAT_WIDTH]), int(stats[label, cv2.CC_STAT_HEIGHT])
        x1, y1 = max(0, x - padding), max(0, y - padding)
        x2, y2 = min(width, x + w + padding), min(height, y + h + padding)
        crop = image[y1:y2, x1:x2].copy()
        crop[labels[y1:y2, x1:x2] != label] = (0, 0, 0, 0)
        name = f"part_{len(saved):02d}.png"
        path = os.path.join(output_dir, name)
        if not cv2.imwrite(path, crop):
            raise ValueError(f"无法写入拆件：{path}")
        saved.append(path)
        components.append({"name": name[:-4], "x": x, "y": y, "width": w, "height": h, "area": area})
    if not saved:
        raise ValueError("拆件图集没有检测到前景，请检查白底和 bg-threshold")
    with open(os.path.join(output_dir, "parts-manifest.json"), "w", encoding="utf-8") as handle:
        json.dump({"atlas_width": width, "atlas_height": height, "parts": components}, handle, ensure_ascii=False, indent=2)
    return saved


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--atlas", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--min-area", type=int, default=500)
    parser.add_argument("--padding", type=int, default=12)
    parser.add_argument("--bg-threshold", type=int, default=240)
    args = parser.parse_args()
    try:
        paths = segment(args.atlas, args.output_dir, args.min_area, args.padding, args.bg_threshold)
    except Exception as error:
        print(f"ERROR: {error}", file=sys.stderr)
        return 1
    print(f"拆出 {len(paths)} 个部件")
    return 0


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