"""Offline packaging of Jackbox / Jukebox art for the MASM32 assembly renderer.

Processes high-resolution source images, removes studio grey background with clean
alpha masking and anti-aliasing, aligns 3D cabinet surfaces (Front, Side, Back),
and outputs binary packs (ANM1).
"""
import math
from pathlib import Path
import struct
import cv2
import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parents[2]
SOURCE = ROOT / "assets/sprites/jackbox"
OUT = ROOT / "assets/animations/hall"


def clean_cutout(img_bgr: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Segment studio grey background using connected components from image borders."""
    h, w = img_bgr.shape[:2]
    arr = img_bgr.astype(np.float32)

    # Detect neutral studio grey pixels: low saturation and mid-low brightness
    color_diff = np.maximum(
        np.abs(arr[:, :, 0] - arr[:, :, 1]),
        np.maximum(np.abs(arr[:, :, 1] - arr[:, :, 2]), np.abs(arr[:, :, 0] - arr[:, :, 2]))
    )
    brightness = np.mean(arr, axis=2)
    is_bg_cand = (color_diff < 10) & (brightness >= 18) & (brightness <= 65)

    # Seed from all 4 borders
    seed_mask = np.zeros((h, w), dtype=np.uint8)
    seed_mask[0, :] = is_bg_cand[0, :]
    seed_mask[-1, :] = is_bg_cand[-1, :]
    seed_mask[:, 0] = is_bg_cand[:, 0]
    seed_mask[:, -1] = is_bg_cand[:, -1]

    num_labels, labels = cv2.connectedComponents(is_bg_cand.astype(np.uint8))
    border_labels = np.unique(labels[seed_mask == 1])
    border_labels = border_labels[border_labels != 0]

    is_bg = np.isin(labels, border_labels)
    clean_fg = (~is_bg).astype(np.uint8)

    # Fill internal holes in foreground if any
    num_fg, fg_labels = cv2.connectedComponents(clean_fg)
    fg_sizes = np.bincount(fg_labels.flatten())
    fg_sizes[0] = 0
    main_fg_label = np.argmax(fg_sizes)
    clean_fg = (fg_labels == main_fg_label).astype(np.uint8)

    # Anti-aliased alpha falloff via distance transform
    dist_in = cv2.distanceTransform(clean_fg, cv2.DIST_L2, 3)
    dist_out = cv2.distanceTransform(1 - clean_fg, cv2.DIST_L2, 3)
    alpha = np.clip((dist_in - dist_out + 0.5) * 255.0, 0, 255).astype(np.uint8)

    bgra = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2BGRA)
    bgra[:, :, 3] = alpha
    return bgra, clean_fg


def pack_bgra(path: Path, frames: list[np.ndarray], w: int, h: int):
    """Pack an array of (H, W, 4) BGRA numpy uint8 images into an ANM1 binary pack."""
    raw = b"".join(f.tobytes() for f in frames)
    path.parent.mkdir(parents=True, exist_ok=True)
    header = struct.pack("<4s4I", b"ANM1", w, h, len(frames), len(raw))
    path.write_bytes(header + raw)
    print(f"Packed {len(frames)} frame(s) ({w}x{h}) into {path.name}: {len(raw)} bytes")


def main():
    front_path = SOURCE / "Jackbox_front.jpg"
    side_path = SOURCE / "Jackbox_side.jpg"
    ui_path = SOURCE / "Jackbox_UI.jpg"

    assert front_path.exists(), f"Missing {front_path}"
    assert side_path.exists(), f"Missing {side_path}"
    assert ui_path.exists(), f"Missing {ui_path}"

    print("Extracting and segmenting Jukebox textures...")

    # 1. Front Texture
    front_bgr = cv2.imread(str(front_path))
    front_bgra, front_fg = clean_cutout(front_bgr)
    front_tex = cv2.resize(front_bgra, (256, 512), interpolation=cv2.INTER_LANCZOS4)

    # 2. Side Texture: clean, crop margins, flip horizontally so back is x=0 and front is x=255
    side_bgr = cv2.imread(str(side_path))
    side_bgra, side_fg = clean_cutout(side_bgr)
    # Crop to actual jukebox silhouette (remove studio floor and top empty margins)
    side_cropped = side_bgra[80:2912, 0:1350]
    side_flipped = cv2.flip(side_cropped, 1)
    side_tex = cv2.resize(side_flipped, (256, 512), interpolation=cv2.INTER_LANCZOS4)

    # 3. Rear Texture: authentic dark walnut cabinet panel with ventilation slots and arch top
    rear_tex = np.zeros((512, 256, 4), dtype=np.uint8)
    for y in range(512):
        grain = int(math.sin(y * 0.2) * 5 + math.sin(y * 0.05) * 8)
        rear_tex[y, :, 0] = np.clip(20 + grain, 0, 255)  # B
        rear_tex[y, :, 1] = np.clip(28 + grain + 2, 0, 255)  # G
        rear_tex[y, :, 2] = np.clip(45 + grain + 5, 0, 255)  # R
    for y in range(120, 220, 10):
        rear_tex[y:y+3, 40:216, :3] = 12  # Dark ventilation slots
    rear_tex[:, :, 3] = cv2.resize(front_bgra[:, :, 3], (256, 512), interpolation=cv2.INTER_LANCZOS4)

    # Pack 3D world frames: Frame 0 = Front, Frame 1 = Side, Frame 2 = Rear
    pack_bgra(OUT / "jackbox-v1.bin", [front_tex, side_tex, rear_tex], 256, 512)

    # 4. Screen UI Pack: 640x480 BGRA
    ui_bgr = cv2.imread(str(ui_path))
    ui_bgra = cv2.cvtColor(ui_bgr, cv2.COLOR_BGR2BGRA)
    ui_tex = cv2.resize(ui_bgra, (640, 480), interpolation=cv2.INTER_LANCZOS4)
    pack_bgra(OUT / "jackbox-ui-v1.bin", [ui_tex], 640, 480)

    print("Jukebox asset packaging complete.")


if __name__ == "__main__":
    main()