app.py

"""Wardrobe AI — Turn your wardrobe into a searchable knowledge graph.

A Gradio app that uses Gemma 3 4B (GGUF) to extract garment attributes
from photos and answer natural language questions about your clothes.

Built for the Build Small Hackathon (HuggingFace x Gradio, June 2026).
"""

import io
import logging
import os
from pathlib import Path

from dotenv import load_dotenv

load_dotenv(Path(__file__).resolve().parent / ".env")

import gradio as gr
from PIL import Image

from src.vision import extract_garments, extract_single_from_path, extract_from_crop_bytes, _extract_single_garment
from src.catalog import (
    add_garments,
    load_catalog,
    clear_catalog,
    get_catalog_summary,
    get_catalog_stats,
    get_garment_image_path,
)
from src.assistant import ask_streaming
from src.combinations import (
    generate_combinations,
    rank_with_llm,
    save_preference,
    get_liked_outfits,
)
from src.detector import detect_garments as detect_boxes, crop_garments, list_available, BoundingBox
from src import settings
from gradio_image_annotation import image_annotator

if os.environ.get("SPACE_ID"):
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
else:
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
logging.basicConfig(level=logging.INFO, format="%(name)s | %(message)s")

_DATA_DIR = Path(__file__).resolve().parent / "data"
_DATA_DIR.mkdir(parents=True, exist_ok=True)
(_DATA_DIR / "garments").mkdir(parents=True, exist_ok=True)

_startup_catalog = load_catalog()
logging.info("Startup: catalog loaded with %d garments", len(_startup_catalog))
del _startup_catalog

HEADER = """
# 👕 Wardrobe AI

**Mi madre tiene más de 200 prendas.** Cada mañana pierde tiempo buscando
algo que combine. No recuerda qué ropa tiene. Compra ropa duplicada.

Wardrobe AI transforma un armario físico en un catálogo consultable
mediante lenguaje natural. Sube fotos → detecta prendas → pregunta lo que quieras.

<small>Gemma 3 4B · 100% local · sin APIs externas · llama.cpp</small>
"""


# ---------------------------------------------------------------------------
# Capture tab handlers
# ---------------------------------------------------------------------------

def _format_results(added: list[dict]) -> tuple[str, list]:
    """Format detection results as (status_markdown, table_data)."""
    if not added:
        return "No se detectaron prendas.", []

    table_data = [
        [g["id"], g["type"], g["color"], g["material"], g["pattern"], g["season"], g["formality"]]
        for g in added
    ]
    status = f"Se detectaron **{len(added)}** prendas nuevas y se añadieron al catálogo."
    return status, table_data


def auto_detect(annotation, mode, progress=gr.Progress(track_tqdm=False)):
    """Run automatic detection on the uploaded image.

    Uses the configured detection backend to find garments, crops them,
    and sends each crop to the VLM for attribute extraction.
    """
    if not annotation or not annotation.get("image"):
        return "Sube una imagen primero.", []

    image_path = annotation["image"]
    if not isinstance(image_path, str):
        return "Error: imagen no válida.", []

    progress((0, 1), desc="Detectando prendas...")

    try:
        if mode == "single":
            garments_with_crops = extract_single_from_path(image_path)
        else:
            garments_with_crops = extract_garments(image_path)
    except Exception as e:
        logging.error("Error in auto detection: %s", e)
        return f"Error durante la detección: {e}", []

    progress((1, 1), desc="Completado")

    if not garments_with_crops:
        return (
            "**No se detectaron prendas automáticamente.** "
            "Prueba a dibujar rectángulos sobre las prendas y pulsa 'Procesar selección manual'.",
            [],
        )

    added = add_garments(garments_with_crops)
    return _format_results(added)


def process_manual_selection(annotation, progress=gr.Progress(track_tqdm=False)):
    """Process manually drawn bounding boxes from the annotator.

    User draws rectangles over garments, we crop each one and send to VLM.
    """
    if not annotation or "boxes" not in annotation or not annotation["boxes"]:
        return "Dibuja al menos un rectángulo sobre una prenda.", []

    image = annotation["image"]
    boxes = annotation["boxes"]

    if isinstance(image, str):
        img = Image.open(image).convert("RGB")
    else:
        img = Image.fromarray(image).convert("RGB")

    total = len(boxes)
    all_results = []
    for idx, box in enumerate(boxes):
        progress((idx, total), desc=f"Analizando recorte {idx + 1}/{total}...")
        x1 = int(box["xmin"])
        y1 = int(box["ymin"])
        x2 = int(box["xmax"])
        y2 = int(box["ymax"])

        cropped = img.crop((x1, y1, x2, y2))
        cropped.thumbnail((512, 512), Image.LANCZOS)

        buffer = io.BytesIO()
        cropped.save(buffer, format="JPEG", quality=85)
        crop_bytes = buffer.getvalue()

        result = extract_from_crop_bytes(crop_bytes)
        if result:
            all_results.append(result)

    progress((total, total), desc="Completado")

    if not all_results:
        return "No se pudieron analizar las prendas seleccionadas.", []

    added = add_garments(all_results)
    return _format_results(added)


# ---------------------------------------------------------------------------
# Settings handlers
# ---------------------------------------------------------------------------

def update_detection_backend(backend_name: str):
    """Persist detection backend change."""
    settings.update("detection_backend", backend_name)
    return f"Backend actualizado: **{backend_name}**"


# ---------------------------------------------------------------------------
# Dataset loading handler (in-process with real-time progress)
# ---------------------------------------------------------------------------

DATA_DIR = Path(__file__).resolve().parent / "data"

SAMPLE_DATASETS = [
    ("second-hand", "Second-hand (prendas individuales)"),
    ("fashion-1k", "Fashion-1K (multi-garment, requiere detección)"),
]

TARGET_GARMENTS = 50


def obtener_dataset(dataset_key: str):
    """Download a HF dataset and process garments in-process with real-time UI updates.

    Yields (log_markdown, gallery_images) at each step.
    """
    from datasets import load_dataset

    ds_configs = {
        "second-hand": {
            "hf_id": "fnauman/fashion-second-hand-front-only-rgb",
            "needs_detection": False,
        },
        "fashion-1k": {
            "hf_id": "Codatta/Fashion-1K",
            "needs_detection": True,
        },
    }

    config = ds_configs.get(dataset_key)
    if not config:
        yield "Error: dataset no reconocido.", []
        return

    yield f"Descargando dataset **{config['hf_id']}**...", []

    try:
        ds = load_dataset(config["hf_id"], split="train")
    except Exception as e:
        yield f"Error descargando dataset: {e}", []
        return

    yield f"Dataset cargado: **{len(ds)}** imágenes. Iniciando procesamiento...", []

    # Spread indices for variety
    step = max(1, len(ds) // (TARGET_GARMENTS * 2))
    indices = list(range(0, len(ds), step))[:TARGET_GARMENTS * 3]

    gallery_images: list[str] = []
    garments_processed: list[tuple[dict, bytes]] = []
    log_lines: list[str] = []

    for idx in indices:
        if len(garments_processed) >= TARGET_GARMENTS:
            break
        if idx >= len(ds):
            continue

        sample = ds[idx]
        image = sample.get("image") or sample.get("img")
        if not isinstance(image, Image.Image):
            continue
        if image.mode != "RGB":
            image = image.convert("RGB")

        if config["needs_detection"]:
            import tempfile
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
                image.save(tmp, format="JPEG", quality=90)
                tmp_path = tmp.name

            try:
                from src.detector import detect_and_crop
                crops = detect_and_crop(tmp_path)
            except Exception:
                crops = []
            finally:
                Path(tmp_path).unlink(missing_ok=True)

            if not crops:
                continue

            for crop_bytes in crops:
                if len(garments_processed) >= TARGET_GARMENTS:
                    break
                garment = _extract_single_garment(crop_bytes)
                if garment:
                    garments_processed.append((garment, crop_bytes))
                    img_path = _save_temp_preview(crop_bytes, len(garments_processed))
                    if img_path:
                        gallery_images.append(img_path)
                    n = len(garments_processed)
                    log_lines.append(
                        f"**{n}/{TARGET_GARMENTS}** — {garment.get('color', '?')} {garment.get('type', '?')}"
                    )
                    yield "\n".join(log_lines[-8:]), gallery_images
        else:
            if max(image.size) > 512:
                image.thumbnail((512, 512), Image.LANCZOS)

            buf = io.BytesIO()
            image.save(buf, format="JPEG", quality=90)
            crop_bytes = buf.getvalue()

            garment = _extract_single_garment(crop_bytes)
            if not garment:
                continue

            garments_processed.append((garment, crop_bytes))
            img_path = _save_temp_preview(crop_bytes, len(garments_processed))
            if img_path:
                gallery_images.append(img_path)
            n = len(garments_processed)
            log_lines.append(
                f"**{n}/{TARGET_GARMENTS}** — {garment.get('color', '?')} {garment.get('type', '?')}"
            )
            yield "\n".join(log_lines[-8:]), gallery_images

    if not garments_processed:
        yield "No se pudieron extraer prendas del dataset.", gallery_images
        return

    # Save to catalog
    added = add_garments(garments_processed)
    final_log = "\n".join(log_lines[-5:]) + (
        f"\n\n---\n**Completado: {len(added)} prendas** añadidas al armario. "
        f"Ve a 'Mi Armario' para explorarlas."
    )
    yield final_log, gallery_images


def _save_temp_preview(crop_bytes: bytes, index: int) -> str | None:
    """Save a preview image for the gallery during dataset loading."""
    preview_dir = DATA_DIR / "garments"
    preview_dir.mkdir(parents=True, exist_ok=True)
    path = preview_dir / f"_preview_{index:03d}.jpg"
    try:
        path.write_bytes(crop_bytes)
        return str(path)
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Wardrobe tab handlers
# ---------------------------------------------------------------------------

def _get_catalog_choices() -> list[tuple[str, str]]:
    """Build (label, garment_id) pairs for the garment list."""
    catalog = load_catalog()
    choices = []
    for g in catalog:
        gid = g.get("id", "?")
        desc = g.get("description", "")
        if desc:
            label = desc[:80]
        else:
            color = g.get("color", "?").title()
            gtype = g.get("type", "?")
            pattern = g.get("pattern", "solid")
            label = f"{color} {gtype} ({pattern})"
        choices.append((label, gid))
    return choices


def _format_garment_detail(garment_id: str) -> tuple[str | None, str]:
    """Return (image_path, detail_markdown) for a garment."""
    catalog = load_catalog()
    garment = next((g for g in catalog if g.get("id") == garment_id), None)

    if not garment:
        return None, "Selecciona una prenda de la lista."

    img_path = get_garment_image_path(garment_id)

    lines = [f"### {garment.get('color', '?').title()} {garment.get('type', '?').title()}"]

    desc = garment.get("description", "")
    if desc:
        lines.append(f"\n*{desc}*\n")

    if not img_path:
        lines.append("\n*Sin imagen disponible*\n")

    lines.append("\n| Atributo | Valor |")
    lines.append("|----------|-------|")
    lines.append(f"| **ID** | `{garment.get('id', '?')}` |")
    lines.append(f"| **Tipo** | {garment.get('type', '?')} |")
    lines.append(f"| **Color** | {garment.get('color', '?')} |")
    lines.append(f"| **Material** | {garment.get('material', '?')} |")
    lines.append(f"| **Patrón** | {garment.get('pattern', '?')} |")
    lines.append(f"| **Temporada** | {garment.get('season', '?')} |")
    lines.append(f"| **Estilo** | {garment.get('formality', '?')} |")

    return img_path, "\n".join(lines)


def _format_catalog_stats() -> str:
    """Format aggregate stats for display above the list."""
    stats = get_catalog_stats()
    if stats["total"] == 0:
        return "El catálogo está vacío. Sube fotos de tu ropa para empezar."

    lines = [f"**{stats['total']} prendas** en tu armario"]

    if stats.get("by_type"):
        type_summary = ", ".join(f"{v}x {k}" for k, v in list(stats["by_type"].items())[:6])
        lines.append(f"Tipos: {type_summary}")

    return " · ".join(lines)


def select_garment(garment_id: str):
    """Handle garment selection from the list."""
    if not garment_id:
        return gr.update(value=None), "Selecciona una prenda de la lista."
    img_path, detail = _format_garment_detail(garment_id)
    return gr.update(value=img_path), detail


def refresh_catalog():
    """Refresh the catalog list and stats."""
    choices = _get_catalog_choices()
    stats = _format_catalog_stats()
    return (
        gr.update(choices=choices, value=None),
        stats,
        gr.update(value=None),
        "Selecciona una prenda de la lista.",
    )


def clear_all():
    """Clear the entire catalog."""
    clear_catalog()
    return (
        gr.update(choices=[], value=None),
        "Catálogo vaciado.",
        gr.update(value=None),
        "El catálogo está vacío.",
    )


# ---------------------------------------------------------------------------
# Chat handlers
# ---------------------------------------------------------------------------

def chat_respond(message, history):
    """Handle chat messages with streaming responses."""
    if not message:
        return ""

    partial = ""
    for token in ask_streaming(message):
        partial += token
        yield partial


EXAMPLE_QUESTIONS = [
    "¿Qué me pongo para una cena informal?",
    "¿Qué combinaciones puedo hacer con los jeans?",
    "¿Cómo debo lavar el sweater de lana?",
    "¿Qué ropa tengo para otoño?",
    "¿Qué me pongo si hace 32 grados?",
    "What should I wear for a job interview?",
]


# ---------------------------------------------------------------------------
# Combinations tab handlers
# ---------------------------------------------------------------------------

def _get_combo_display(combo: dict | None) -> tuple[str | None, str | None, str, str]:
    """Extract display data from a combination."""
    if not combo:
        return None, None, "No hay combinaciones disponibles.", ""

    top = combo["top"]
    bottom = combo["bottom"]

    top_img = get_garment_image_path(top["id"])
    bottom_img = get_garment_image_path(bottom["id"])

    top_text = top.get("description") or f"{top['color'].title()} {top['type']}"
    bottom_text = bottom.get("description") or f"{bottom['color'].title()} {bottom['type']}"

    return top_img, bottom_img, top_text, bottom_text


def init_combinations(state, context):
    """Initialize or refresh the combination queue, optionally ranked by context."""
    combos = generate_combinations()
    if not combos:
        return (
            state,
            None, None,
            "No hay suficientes prendas (necesitas al menos 1 top y 1 bottom).",
            "",
            "0 combinaciones",
        )

    combos = rank_with_llm(combos, context.strip() if context else "")

    state = {"queue": combos, "index": 0}
    combo = combos[0]
    top_img, bottom_img, top_text, bottom_text = _get_combo_display(combo)
    progress = f"1 / {len(combos)}"

    return state, top_img, bottom_img, top_text, bottom_text, progress


def handle_preference(liked: bool, state):
    """Record preference and advance to the next combination."""
    if not state or "queue" not in state:
        return state, None, None, "Genera combinaciones primero.", "", "—"

    queue = state["queue"]
    idx = state["index"]

    if idx < len(queue):
        combo = queue[idx]
        save_preference(combo["top"]["id"], combo["bottom"]["id"], liked)

    idx += 1
    state["index"] = idx

    if idx >= len(queue):
        return (
            state, None, None,
            "Has revisado todas las combinaciones.",
            "Pulsa 'Generar' para crear nuevas.",
            f"{idx} / {len(queue)}",
        )

    combo = queue[idx]
    top_img, bottom_img, top_text, bottom_text = _get_combo_display(combo)
    progress = f"{idx + 1} / {len(queue)}"

    return state, top_img, bottom_img, top_text, bottom_text, progress


def handle_like(state):
    return handle_preference(True, state)


def handle_dislike(state):
    return handle_preference(False, state)


def show_liked_outfits():
    """Format liked outfits for display."""
    liked = get_liked_outfits()
    if not liked:
        return "Aún no has guardado ningún outfit. Usa la tab Combina para aprobar looks."

    lines = [f"### Outfits guardados: {len(liked)}\n"]
    for outfit in liked:
        top = outfit["top"]
        bottom = outfit["bottom"]
        top_desc = top.get("description") or f"{top['color']} {top['type']}"
        bottom_desc = bottom.get("description") or f"{bottom['color']} {bottom['type']}"
        lines.append(f"- **{outfit['id']}**: {top_desc} + {bottom_desc}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# UI Layout
# ---------------------------------------------------------------------------

with gr.Blocks(title="Wardrobe AI") as demo:

    gr.Markdown(HEADER)

    with gr.Tab("📸 Captura"):
        gr.Markdown(
            "Sube una foto de tu ropa. Puedes detectar prendas automáticamente "
            "o dibujar rectángulos sobre ellas para seleccionarlas manualmente."
        )

        with gr.Row():
            with gr.Column(scale=1, min_width=200):
                capture_mode = gr.Radio(
                    choices=[
                        ("Prenda individual", "single"),
                        ("Múltiples prendas", "multi"),
                    ],
                    value="multi",
                    label="Modo de captura",
                )

        bbox_annotator = image_annotator(
            value=None,
            image_type="filepath",
            label_list=["prenda"],
            label="Sube una foto y dibuja rectángulos sobre las prendas (opcional)",
        )

        with gr.Row():
            auto_detect_btn = gr.Button(
                "Detectar automáticamente", variant="primary", size="lg", scale=2,
            )
            manual_btn = gr.Button(
                "Procesar selección manual", variant="secondary", size="lg", scale=2,
            )

        status_output = gr.Markdown(label="Estado")
        detected_table = gr.Dataframe(
            headers=["ID", "Tipo", "Color", "Material", "Patrón", "Temporada", "Estilo"],
            label="Prendas detectadas",
            interactive=False,
        )

        gr.Markdown("---")
        gr.Markdown("#### Obtener dataset de ejemplo")
        with gr.Row():
            dataset_select = gr.Dropdown(
                choices=[(label, key) for key, label in SAMPLE_DATASETS],
                value="second-hand",
                label="Dataset de origen",
                scale=2,
            )
            obtener_btn = gr.Button(
                "Obtener Dataset", variant="primary", size="lg", scale=1,
            )

        dataset_log = gr.Markdown()
        dataset_gallery = gr.Gallery(
            label="Prendas procesadas",
            columns=6,
            rows=2,
            height="auto",
            object_fit="contain",
        )

        obtener_btn.click(
            obtener_dataset,
            inputs=[dataset_select],
            outputs=[dataset_log, dataset_gallery],
        )

        with gr.Accordion("⚙️ Ajustes de detección", open=False):
            available_backends = list_available()
            if not available_backends:
                available_backends = [("YOLOS-tiny (HuggingFace)", "yolos")]

            current_backend = settings.get("detection_backend")
            backend_selector = gr.Radio(
                choices=available_backends,
                value=current_backend,
                label="Backend de detección automática",
            )
            backend_status = gr.Markdown(
                value=f"Backend activo: **{current_backend}**"
            )
            backend_selector.change(
                update_detection_backend,
                inputs=[backend_selector],
                outputs=[backend_status],
            )

        auto_detect_btn.click(
            auto_detect,
            inputs=[bbox_annotator, capture_mode],
            outputs=[status_output, detected_table],
        )
        manual_btn.click(
            process_manual_selection,
            inputs=[bbox_annotator],
            outputs=[status_output, detected_table],
        )

    with gr.Tab("👗 Mi Armario"):
        gr.HTML("""<style>
            .garment-list { max-height: 420px; overflow-y: auto; }
            .garment-list .wrap { gap: 4px !important; }
        </style>""")

        catalog_stats = gr.Markdown(value=_format_catalog_stats())

        with gr.Row():
            refresh_btn = gr.Button("Actualizar", size="sm")
            clear_btn = gr.Button("Vaciar catálogo", variant="stop", size="sm")

        with gr.Row(equal_height=True):
            with gr.Column(scale=1, min_width=220):
                garment_list = gr.Radio(
                    choices=_get_catalog_choices(),
                    label="Prendas",
                    interactive=True,
                    elem_classes=["garment-list"],
                )

            with gr.Column(scale=2, min_width=400):
                with gr.Row(equal_height=True):
                    garment_image = gr.Image(
                        label="Foto",
                        type="filepath",
                        interactive=False,
                        show_label=False,
                        height=280,
                    )
                    garment_detail = gr.Markdown(
                        value="Selecciona una prenda de la lista."
                    )

        garment_list.change(
            select_garment,
            inputs=[garment_list],
            outputs=[garment_image, garment_detail],
        )
        refresh_btn.click(
            refresh_catalog,
            outputs=[garment_list, catalog_stats, garment_image, garment_detail],
        )
        clear_btn.click(
            clear_all,
            outputs=[garment_list, catalog_stats, garment_image, garment_detail],
        )

    with gr.Tab("💬 Pregunta"):
        gr.Markdown("Pregúntale a tu armario. Responde en el idioma de tu pregunta.")

        chatbot = gr.ChatInterface(
            fn=chat_respond,
            examples=EXAMPLE_QUESTIONS,
            title=None,
        )

    with gr.Tab("💫 Combina"):
        gr.Markdown(
            "Desliza entre combinaciones de outfits. "
            "**Top + Bottom** — ¿te gusta o no?"
        )

        combo_state = gr.State(value=None)

        event_context = gr.Textbox(
            placeholder="Ej: cena informal en terraza, 30 grados...",
            label="Describe la ocasión (opcional — las combinaciones se ordenarán para este contexto)",
            lines=1,
        )

        with gr.Row():
            generate_btn = gr.Button(
                "Generar combinaciones", variant="primary", size="lg",
            )
            combo_progress = gr.Markdown(value="—")

        with gr.Row(equal_height=True):
            with gr.Column(scale=1, min_width=200):
                gr.Markdown("#### Top")
                combo_top_img = gr.Image(
                    label="Top",
                    type="filepath",
                    interactive=False,
                    show_label=False,
                    height=280,
                )
                combo_top_text = gr.Markdown()

            with gr.Column(scale=1, min_width=200):
                gr.Markdown("#### Bottom")
                combo_bottom_img = gr.Image(
                    label="Bottom",
                    type="filepath",
                    interactive=False,
                    show_label=False,
                    height=280,
                )
                combo_bottom_text = gr.Markdown()

        with gr.Row():
            dislike_btn = gr.Button("👎 No me gusta", variant="stop", size="lg", scale=1)
            like_btn = gr.Button("👍 Me gusta", variant="primary", size="lg", scale=1)

        gr.Markdown("---")
        liked_display = gr.Markdown(value=show_liked_outfits)
        refresh_liked_btn = gr.Button("Ver outfits guardados", size="sm")

        combo_outputs = [
            combo_state, combo_top_img, combo_bottom_img,
            combo_top_text, combo_bottom_text, combo_progress,
        ]

        generate_btn.click(
            init_combinations,
            inputs=[combo_state, event_context],
            outputs=combo_outputs,
        )
        like_btn.click(
            handle_like,
            inputs=[combo_state],
            outputs=combo_outputs,
        )
        dislike_btn.click(
            handle_dislike,
            inputs=[combo_state],
            outputs=combo_outputs,
        )
        refresh_liked_btn.click(show_liked_outfits, outputs=[liked_display])

    with gr.Accordion("Acerca del proyecto", open=False):
        gr.Markdown(
            """
**Wardrobe AI** transforma un armario físico en un catálogo digital consultable.

- **Modelo**: Gemma 3 4B IT (Q4_K_M GGUF) — 4B parámetros
- **Inferencia**: llama.cpp via llama-cpp-python — 100% local
- **Track**: Backyard AI — "Resuelve un problema real para alguien que conoces"
- **Badges**: Tiny Titan (≤4B) · Off the Grid (sin APIs) · Llama Champion (llama.cpp)

Construido para el [Build Small Hackathon](https://huggingface.co/build-small-hackathon)
(HuggingFace × Gradio, junio 2026).
"""
        )


# ---------------------------------------------------------------------------
# Custom frontend (gr.Server + Alpine.js)
# ---------------------------------------------------------------------------

def _build_custom_server():
    """Build the gr.Server app with API endpoints and custom HTML frontend."""
    from gradio import Server
    from fastapi.responses import FileResponse
    from fastapi.staticfiles import StaticFiles
    from src.assistant import ask
    from src.storage import get_image_path

    UI_DIR = Path(__file__).resolve().parent / "src" / "ui"
    server = Server()

    server.mount("/garments", StaticFiles(directory=str(_DATA_DIR / "garments")), name="garments")

    _UPLOADS_DIR = _DATA_DIR / "_uploads"
    _UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
    server.mount("/uploads", StaticFiles(directory=str(_UPLOADS_DIR)), name="uploads")

    def _image_url(garment_id: str) -> str:
        """Return a cache-busted URL for a garment image using file mtime."""
        img_path = get_image_path(garment_id)
        if img_path:
            try:
                v = int(os.path.getmtime(img_path))
            except OSError:
                v = 0
        else:
            v = 0
        return f"/garments/{garment_id}.jpg?v={v}"

    @server.api(name="prepare_image")
    def api_prepare_image(image_path: str | dict) -> dict:
        """Save image to _uploads, run auto-detection, return token + detected boxes.

        The frontend uses the token + image_url to show the image in Annotorious
        pre-populated with auto-detected boxes for the user to review/edit.
        """
        import uuid
        import time as _time

        if isinstance(image_path, dict):
            image_path = image_path.get("path") or image_path.get("url", "")

        try:
            img = Image.open(str(image_path)).convert("RGB")
        except Exception as e:
            return {"error": str(e), "token": "", "image_url": "", "width": 0, "height": 0, "boxes": []}

        img.thumbnail((1280, 1280), Image.LANCZOS)
        w, h = img.size

        token = uuid.uuid4().hex
        upload_path = _UPLOADS_DIR / f"{token}.jpg"
        img.save(str(upload_path), format="JPEG", quality=90)

        try:
            boxes = detect_boxes(str(upload_path))
        except Exception:
            boxes = []

        ts = int(_time.time())
        return {
            "token": token,
            "image_url": f"/uploads/{token}.jpg?v={ts}",
            "width": w,
            "height": h,
            "boxes": [{"x": b.x1, "y": b.y1, "w": b.width, "h": b.height} for b in boxes],
        }

    @server.api(name="analyze_boxes")
    def api_analyze_boxes(token: str, boxes: str) -> dict:
        """Crop each user-confirmed bounding box from the uploaded image and extract garment attributes.

        Args:
            token: filename token returned by prepare_image (hex, no path separators).
            boxes: JSON string — list of {x, y, w, h} in pixels of the stored image.
        """
        import json as _json
        import re as _re

        # Validate token: only hex characters, no path traversal
        if not _re.fullmatch(r"[0-9a-f]{32}", token):
            return {"error": "Invalid token", "count": 0, "garments": []}

        upload_path = _UPLOADS_DIR / f"{token}.jpg"
        if not upload_path.exists():
            return {"error": "Image not found. Please re-upload.", "count": 0, "garments": []}

        try:
            box_list = _json.loads(boxes)
        except Exception:
            return {"error": "Invalid boxes format", "count": 0, "garments": []}

        try:
            img = Image.open(str(upload_path)).convert("RGB")
        except Exception as e:
            return {"error": str(e), "count": 0, "garments": []}

        results: list[tuple[dict, bytes]] = []
        for box in box_list:
            try:
                x = int(box["x"])
                y = int(box["y"])
                w = int(box["w"])
                h = int(box["h"])
            except (KeyError, TypeError, ValueError):
                continue

            if w < 10 or h < 10:
                continue

            cropped = img.crop((x, y, x + w, y + h))
            cropped.thumbnail((512, 512), Image.LANCZOS)

            buf = io.BytesIO()
            cropped.save(buf, format="JPEG", quality=85)
            crop_bytes = buf.getvalue()

            result = extract_from_crop_bytes(crop_bytes)
            if result:
                results.append(result)

        # Cleanup upload temp file
        try:
            upload_path.unlink()
        except OSError:
            pass

        if not results:
            return {"error": "No garments could be extracted from the selections.", "count": 0, "garments": []}

        added = add_garments(results)
        for g in added:
            g["image_url"] = _image_url(g["id"])

        return {"count": len(added), "garments": added}

    @server.api(name="get_wardrobe")
    def api_get_wardrobe() -> dict:
        catalog = load_catalog()
        for g in catalog:
            g["image_url"] = _image_url(g.get("id", ""))
        return {"garments": catalog, "count": len(catalog)}

    @server.api(name="add_photo")
    def api_add_photo(image_path: str | dict) -> dict:
        if isinstance(image_path, dict):
            image_path = image_path.get("path") or image_path.get("url", "")
        results = extract_garments(str(image_path))
        if not results:
            return {"garments": [], "count": 0}
        added = add_garments(results)
        for g in added:
            g["image_url"] = _image_url(g["id"])
        return {"garments": added, "count": len(added)}

    @server.api(name="get_combinations")
    def api_get_combinations(context: str = "") -> dict:
        combos = generate_combinations()
        if not combos:
            return {"combinations": [], "count": 0}
        combos = rank_with_llm(combos, context.strip() if context else "")
        serialized = []
        for combo in combos[:20]:
            top, bottom = combo["top"], combo["bottom"]
            serialized.append({
                "id": combo["id"],
                "top": {"id": top.get("id", ""), "type": top.get("type", ""),
                        "color": top.get("color", ""), "image_url": _image_url(top.get("id", ""))},
                "bottom": {"id": bottom.get("id", ""), "type": bottom.get("type", ""),
                           "color": bottom.get("color", ""), "image_url": _image_url(bottom.get("id", ""))},
            })
        return {"combinations": serialized, "count": len(serialized)}

    @server.api(name="rate_outfit")
    def api_rate_outfit(top_id: str, bottom_id: str, liked: bool) -> dict:
        save_preference(top_id, bottom_id, liked)
        return {"status": "ok", "liked": liked}

    @server.api(name="ask_question")
    def api_ask_question(question: str) -> str:
        if not question or not question.strip():
            return "Please ask a question about your wardrobe."
        return ask(question.strip())

    @server.api(name="load_dataset")
    def api_load_dataset(dataset_key: str) -> dict:
        """Stream dataset processing progress as each garment is analyzed.

        Yields progress dicts with: done, count, log[], preview_url.
        The final yield has done=True with the total count.
        """
        import time as _time
        from datasets import load_dataset as hf_load

        _load_ts = int(_time.time())

        ds_configs = {
            "second-hand": {"hf_id": "fnauman/fashion-second-hand-front-only-rgb", "needs_detection": False},
            "fashion-1k": {"hf_id": "Codatta/Fashion-1K", "needs_detection": True},
        }
        config = ds_configs.get(dataset_key)
        if not config:
            yield {"done": True, "error": "Unknown dataset", "count": 0, "log": ["Error: dataset not recognized."], "preview_url": None}
            return

        log_lines: list[str] = [f"Downloading {config['hf_id']}..."]
        yield {"done": False, "count": 0, "log": log_lines[:], "preview_url": None}

        try:
            ds = hf_load(config["hf_id"], split="train")
        except Exception as e:
            yield {"done": True, "error": str(e), "count": 0, "log": [f"Error downloading dataset: {e}"], "preview_url": None}
            return

        log_lines.append(f"Dataset loaded: {len(ds)} images. Starting processing...")
        yield {"done": False, "count": 0, "log": log_lines[:], "preview_url": None}

        target = TARGET_GARMENTS
        step = max(1, len(ds) // (target * 2))
        indices = list(range(0, len(ds), step))[:target * 3]
        garments_processed: list[tuple[dict, bytes]] = []

        for idx in indices:
            if len(garments_processed) >= target:
                break
            if idx >= len(ds):
                continue
            sample = ds[idx]
            image = sample.get("image") or sample.get("img")
            if not isinstance(image, Image.Image):
                continue
            if image.mode != "RGB":
                image = image.convert("RGB")

            if config["needs_detection"]:
                import tempfile
                with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
                    image.save(tmp, format="JPEG", quality=90)
                    tmp_path = tmp.name
                try:
                    from src.detector import detect_and_crop
                    crops = detect_and_crop(tmp_path)
                except Exception:
                    crops = []
                finally:
                    Path(tmp_path).unlink(missing_ok=True)
                for crop_bytes in crops:
                    if len(garments_processed) >= target:
                        break
                    garment = _extract_single_garment(crop_bytes)
                    if garment:
                        garments_processed.append((garment, crop_bytes))
                        n = len(garments_processed)
                        preview_path = _save_temp_preview(crop_bytes, n)
                        preview_url = f"/garments/_preview_{n:03d}.jpg?v={_load_ts}" if preview_path else None
                        log_lines.append(f"{n}/{target} — {garment.get('color', '?')} {garment.get('type', '?')}")
                        yield {"done": False, "count": n, "log": log_lines[-12:], "preview_url": preview_url}
            else:
                if max(image.size) > 512:
                    image.thumbnail((512, 512), Image.LANCZOS)
                buf = io.BytesIO()
                image.save(buf, format="JPEG", quality=90)
                crop_bytes = buf.getvalue()
                garment = _extract_single_garment(crop_bytes)
                if garment:
                    garments_processed.append((garment, crop_bytes))
                    n = len(garments_processed)
                    preview_path = _save_temp_preview(crop_bytes, n)
                    preview_url = f"/garments/_preview_{n:03d}.jpg?v={_load_ts}" if preview_path else None
                    log_lines.append(f"{n}/{target} — {garment.get('color', '?')} {garment.get('type', '?')}")
                    yield {"done": False, "count": n, "log": log_lines[-12:], "preview_url": preview_url}

        if not garments_processed:
            yield {"done": True, "error": "No garments extracted", "count": 0, "log": log_lines + ["No garments could be extracted."], "preview_url": None}
            return

        added = add_garments(garments_processed)
        log_lines.append(f"Done: {len(added)} garments added to your wardrobe.")
        yield {"done": True, "count": len(added), "log": log_lines[-12:], "preview_url": None}

    @server.get("/")
    async def homepage():
        return FileResponse(str(UI_DIR / "index.html"), media_type="text/html")

    @server.get("/style.css")
    async def styles():
        return FileResponse(str(UI_DIR / "style.css"), media_type="text/css")

    return server


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Wardrobe AI")
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--ui", action="store_true", default=True,
        help="Launch custom frontend (default)",
    )
    group.add_argument(
        "--default", action="store_true",
        help="Launch standard Gradio Blocks UI",
    )
    args = parser.parse_args()

    is_spaces = os.environ.get("SPACE_ID") is not None
    server_name = "0.0.0.0" if is_spaces else "127.0.0.1"

    if args.default:
        demo.launch(
            server_name=server_name,
            server_port=7860,
            theme=gr.themes.Soft(
                primary_hue="stone",
                secondary_hue="amber",
            ),
        )
    else:
        server = _build_custom_server()
        server.launch(
            server_name=server_name,
            server_port=7860,
        )