#!/usr/bin/env python3
"""
Kiwi logo compositor.

Takes:
  - A kiwifruit disc SVG (generated by generate_disc.py)
  - A kiwi bird PNG (AI-generated via ComfyUI img2img then PIL-corrected)

Composites bird over disc, removes white background, outputs icon sizes.

Usage:
    conda run -n cf python3 scripts/logo/composite.py \
        --disc scripts/logo/kiwi_disc.svg \
        --bird scripts/logo/kiwi_bird_final.png \
        --out frontend/public/icons/

Requirements:
    cairosvg, Pillow, numpy
"""

import argparse
import io
from pathlib import Path

import cairosvg
import numpy as np
from PIL import Image, ImageDraw

SIZES = [512, 256, 192, 64]


def remove_background(bird_img: Image.Image) -> Image.Image:
    """Remove white/gray background from AI-generated bird PNG."""
    arr = np.array(bird_img, dtype=np.float32)
    r, g, b = arr[:, :, 0], arr[:, :, 1], arr[:, :, 2]

    brightness = (r + g + b) / 3.0
    saturation = np.max([r, g, b], axis=0) - np.min([r, g, b], axis=0)

    near_white = (brightness > 228) & (saturation < 28)

    # Remove gray ground shadow in bottom 30% of image only
    h = arr.shape[0]
    gray_shadow = np.zeros(arr.shape[:2], dtype=bool)
    gray_shadow[int(h * 0.70):, :] = (
        (brightness[int(h * 0.70):, :] > 135) &
        (saturation[int(h * 0.70):, :] < 65)
    )
    gray_shadow2 = np.zeros(arr.shape[:2], dtype=bool)
    gray_shadow2[int(h * 0.78):, :] = (
        (brightness[int(h * 0.78):, :] > 110) &
        (saturation[int(h * 0.78):, :] < 75)
    )

    transparent = near_white | gray_shadow | gray_shadow2

    # Soft edge
    soft = (brightness > 205) & (saturation < 55) & ~transparent
    soft_alpha = np.where(soft, ((255.0 - brightness) / 50.0) * 255.0, 0.0)

    new_alpha = np.where(transparent, 0.0, np.where(soft, soft_alpha, 255.0))
    arr[:, :, 3] = np.clip(new_alpha, 0, 255)
    return Image.fromarray(arr.astype(np.uint8))


def composite(disc_svg: Path, bird_png: Path, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)

    # Render disc at max size
    max_size = max(SIZES)
    disc_bytes = cairosvg.svg2png(url=str(disc_svg), output_width=max_size, output_height=max_size)
    disc = Image.open(io.BytesIO(disc_bytes)).convert("RGBA")

    # Load and de-background bird
    bird_raw = Image.open(bird_png).convert("RGBA")
    bird_nobg = remove_background(bird_raw)
    bbox = bird_nobg.getbbox()
    bird_crop = bird_nobg.crop(bbox)

    # Scale bird to 80% of disc diameter
    target_w = int(max_size * 0.80)
    target_h = int(bird_crop.height * target_w / bird_crop.width)
    bird_scaled = bird_crop.resize((target_w, target_h), Image.LANCZOS)

    # Composite centered, slightly raised and shifted right (beak faces left)
    cx, cy = max_size // 2, max_size // 2
    x_off = cx - bird_scaled.width // 2 + int(max_size * 0.06)
    y_off = cy - bird_scaled.height // 2 - int(max_size * 0.04)

    result = disc.copy()
    result.paste(bird_scaled, (x_off, y_off), bird_scaled)

    # Draw eye post-composite (eye was at 305,335 in 768px bird space)
    # Eye was at pixel (305, 335) in original 768px bird image
    # Adjust for crop bbox offset, then scale to composite space
    sx = target_w / bird_crop.width
    sy = target_h / bird_crop.height
    # Known eye coords from PIL correction pass (EX=305, EY=335 in 768px space)
    eye_x = int(x_off + (305 - bbox[0]) * sx)
    eye_y = int(y_off + (335 - bbox[1]) * sy)
    draw = ImageDraw.Draw(result)
    draw.ellipse([eye_x - 7, eye_y - 7, eye_x + 7, eye_y + 7], fill=(240, 235, 220, 255))
    draw.ellipse([eye_x - 3, eye_y - 3, eye_x + 3, eye_y + 3], fill=(12, 5, 2, 255))

    # Export all sizes
    for size in SIZES:
        out_path = out_dir / f"icon-{size}.png"
        resized = result.resize((size, size), Image.LANCZOS)
        resized.save(out_path)
        print(f"  Saved {out_path} ({size}px)")

    # Maskable copies (same image — kiwi disc has safe-zone padding built in)
    for size in [192, 512]:
        src = out_dir / f"icon-{size}.png"
        dst = out_dir / f"maskable-{size}.png"
        dst.write_bytes(src.read_bytes())
        print(f"  Copied {dst}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Kiwi logo compositor")
    parser.add_argument("--disc", type=Path, default=Path("scripts/logo/kiwi_disc.svg"))
    parser.add_argument("--bird", type=Path, default=Path("scripts/logo/kiwi_bird_final.png"))
    parser.add_argument("--out", type=Path, default=Path("frontend/public/icons"))
    args = parser.parse_args()

    print(f"Compositing {args.bird.name} onto {args.disc.name} → {args.out}/")
    composite(args.disc, args.bird, args.out)
    print("Done.")