#!/usr/bin/env python3
"""Build prototype/assets/garments/photo: the demo wardrobe as real, cut-out garment photos.

Each garment is a free Pexels photo (Pexels License: free for commercial use, no attribution
required). The script crops the photo to the garment, removes the background with rembg
(the same kind of segmentation the real upload pipeline would run on the server), trims it
and saves a transparent WebP.

rembg is heavy, so it runs from its own environment:
    uv venv --python 3.12 /tmp/rembg-env && VIRTUAL_ENV=/tmp/rembg-env uv pip install "rembg[cpu]"
    /tmp/rembg-env/bin/python tools/build_garment_photos.py [--cache DIR]
"""
from __future__ import annotations

import argparse
import tempfile
import urllib.request
from pathlib import Path

from PIL import Image
from rembg import new_session, remove

ROOT = Path(__file__).resolve().parent.parent
OUT_DIR = ROOT / "prototype" / "assets" / "garments" / "photo"
USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64) SKAP-build-garment-photos"
SOURCE_WIDTH = 1100
MAX_SIDE = 640
ALPHA_FLOOR = 24
WEBP_QUALITY = 84
FULL = (0.0, 0.0, 1.0, 1.0)

# garment id: (Pexels photo id, crop box as fractions left, top, right, bottom, rembg model)
GARMENTS = {
    "coat-trench": ("5424917", FULL, "isnet-general-use"),
    "coat-wool": ("4057673", (0.33, 0.10, 0.648, 0.73), "u2net"),
    "jacket-puffer": ("6044973", (0.25, 0.39, 0.86, 0.98), "u2net"),
    "jacket-rain": ("8113002", FULL, "isnet-general-use"),
    "blazer-navy": ("20428090", (0.24, 0.37, 0.82, 0.80), "u2net"),
    "tee-white": ("12025472", FULL, "isnet-general-use"),
    "breton-top": ("12955555", FULL, "isnet-general-use"),
    "oxford-shirt": ("9558723", (0.24, 0.31, 0.73, 0.97), "u2net"),
    "turtleneck-black": ("3353621", (0.0, 0.30, 0.90, 1.0), "u2net"),
    "jeans-straight": ("1082528", (0.50, 0.05, 0.82, 0.83), "u2net"),
    "trousers-black": ("1082528", (0.15, 0.05, 0.45, 0.83), "u2net"),
    "skirt-pleated": ("38266547", FULL, "isnet-general-use"),
    "dress-slip": ("35538618", FULL, "isnet-general-use"),
    "dress-shirt": ("31683080", FULL, "isnet-general-use"),
    "sneakers-white": ("35618111", FULL, "isnet-general-use"),
    "boots-chelsea": ("40140191", FULL, "isnet-general-use"),
    "loafers-brown": ("25381178", FULL, "isnet-general-use"),
    "boots-rubber": ("40028276", FULL, "isnet-general-use"),
    "sunglasses": ("35989302", FULL, "isnet-general-use"),
    "beanie-red": ("11512730", FULL, "isnet-general-use"),
    "bag-tote": ("12025443", FULL, "isnet-general-use"),
    "bag-shoulder": ("22434764", FULL, "isnet-general-use"),
}


def fetch(photo_id: str, cache: Path) -> Path:
    target = cache / f"{photo_id}.jpg"
    if not target.is_file():
        url = f"https://images.pexels.com/photos/{photo_id}/pexels-photo-{photo_id}.jpeg?auto=compress&cs=tinysrgb&w={SOURCE_WIDTH}"
        request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        with urllib.request.urlopen(request, timeout=60) as response:
            target.write_bytes(response.read())
    return target


def cut_out(source: Path, box: tuple[float, float, float, float], session) -> Image.Image:
    image = Image.open(source).convert("RGB")
    width, height = image.size
    image = image.crop((int(box[0] * width), int(box[1] * height), int(box[2] * width), int(box[3] * height)))
    cut = remove(image, session=session, post_process_mask=True)
    bounds = cut.getchannel("A").point(lambda alpha: 255 if alpha > ALPHA_FLOOR else 0).getbbox()
    if bounds:
        cut = cut.crop(bounds)
    cut.thumbnail((MAX_SIDE, MAX_SIDE))
    return cut


def main() -> None:
    parser = argparse.ArgumentParser(description="Cut out the demo wardrobe photos.")
    parser.add_argument("--cache", type=Path, default=Path(tempfile.gettempdir()) / "skap-garment-photos")
    parser.add_argument("only", nargs="*", help="garment ids to rebuild (default: all)")
    args = parser.parse_args()
    args.cache.mkdir(parents=True, exist_ok=True)
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    sessions = {}
    for garment, (photo_id, box, model) in GARMENTS.items():
        if args.only and garment not in args.only:
            continue
        session = sessions.setdefault(model, new_session(model))
        cut = cut_out(fetch(photo_id, args.cache), box, session)
        target = OUT_DIR / f"{garment}.webp"
        cut.save(target, "WEBP", quality=WEBP_QUALITY, method=6)
        print(f"{garment}: {cut.size[0]}x{cut.size[1]}, {target.stat().st_size // 1024} KB")


if __name__ == "__main__":
    main()