src/vision.py

"""Vision pipeline: extract garment attributes from images using a VLM.

The pipeline uses a pluggable detector to locate individual garments in
an image, crops each one, then sends each crop to Gemma 3 4B via
llama-cpp-python for structured attribute extraction. Each garment gets
its own unique thumbnail.
"""

import base64
import io
import json
import logging
import re
from pathlib import Path

from PIL import Image

from .model_loader import model_manager, GARMENT_TYPES
from .detector import detect_and_crop

logger = logging.getLogger(__name__)

MAX_IMAGE_PIXELS = 512

SINGLE_GARMENT_PROMPT = """Analyze this image of a single clothing item carefully. Return a JSON object with these exact fields:
- "type": garment type (e.g. "sweater", "shirt", "jeans", "boots", "hat", "scarf", "belt", "bag")
- "color": primary color (e.g. "red", "blue", "black", "brown", "white", "beige")
- "material": fabric/material (e.g. "knit", "denim", "leather", "cotton", "silk", "polyester"), or "unknown"
- "pattern": pattern type ("solid", "checkered", "striped", "floral", "cable-knit", "plaid"), or "solid"
- "season": best season ("spring", "summer", "autumn", "winter", "all")
- "formality": style level ("casual", "smart-casual", "formal")
- "description": a short natural language description (1-2 sentences) of the garment including its style, fit, and any notable visual details. Example: "Chunky cable-knit oversized sweater in deep red with crew neck, warm and cozy for layering in cold weather."

Return ONLY a valid JSON object. No explanation, no markdown fences."""

MULTI_GARMENT_PROMPT = """Analyze this image of clothing items carefully. For EACH visible garment, shoe, or fashion accessory, return a JSON array of objects.

Each object MUST have these exact fields:
- "type": garment type (e.g. "sweater", "shirt", "jeans", "boots", "hat", "scarf", "belt", "bag")
- "color": primary color (e.g. "red", "blue", "black", "brown", "white", "beige")
- "material": fabric/material (e.g. "knit", "denim", "leather", "cotton", "silk", "polyester"), or "unknown"
- "pattern": pattern type ("solid", "checkered", "striped", "floral", "cable-knit", "plaid"), or "solid"
- "season": best season ("spring", "summer", "autumn", "winter", "all")
- "formality": style level ("casual", "smart-casual", "formal")
- "description": a short natural language description (1-2 sentences) of the garment including its style, fit, and any notable visual details. Example: "Chunky cable-knit oversized sweater in deep red with crew neck, warm and cozy for layering in cold weather."

IMPORTANT: Only include clothing items, shoes, and fashion accessories. Do NOT include cameras, electronics, decorations, or other non-clothing objects.

Return ONLY a valid JSON array. No explanation, no markdown fences."""


def _image_bytes_to_data_uri(jpeg_bytes: bytes) -> str:
    """Convert JPEG bytes to a base64 data URI."""
    b64 = base64.b64encode(jpeg_bytes).decode("utf-8")
    return f"data:image/jpeg;base64,{b64}"


def _prepare_image(image_path: str) -> tuple[str, bytes]:
    """Resize image and convert to base64 data URI.

    Returns (data_uri, jpeg_bytes) so the thumbnail can be persisted.
    """
    img = Image.open(image_path)
    if img.mode == "RGBA":
        img = img.convert("RGB")
    img.thumbnail((MAX_IMAGE_PIXELS, MAX_IMAGE_PIXELS), Image.LANCZOS)

    buffer = io.BytesIO()
    img.save(buffer, format="JPEG", quality=85)
    jpeg_bytes = buffer.getvalue()
    b64 = base64.b64encode(jpeg_bytes).decode("utf-8")
    return f"data:image/jpeg;base64,{b64}", jpeg_bytes


def _parse_json_response(text: str) -> list[dict]:
    """Extract JSON array from model response, handling common formatting issues."""
    cleaned = text.strip()

    fence_match = re.search(r"```(?:json)?\s*\n?(.*?)```", cleaned, re.DOTALL)
    if fence_match:
        cleaned = fence_match.group(1).strip()

    try:
        parsed = json.loads(cleaned)
        if isinstance(parsed, list):
            return parsed
        if isinstance(parsed, dict):
            return [parsed]
    except json.JSONDecodeError:
        pass

    start = cleaned.find("[")
    end = cleaned.rfind("]")
    if start != -1 and end != -1 and end > start:
        try:
            return json.loads(cleaned[start:end + 1])
        except json.JSONDecodeError:
            pass

    obj_start = cleaned.find("{")
    obj_end = cleaned.rfind("}")
    if obj_start != -1 and obj_end != -1 and obj_end > obj_start:
        try:
            parsed = json.loads(cleaned[obj_start:obj_end + 1])
            if isinstance(parsed, dict):
                return [parsed]
        except json.JSONDecodeError:
            pass

    logger.warning("Could not parse JSON from response: %s", cleaned[:200])
    return []


def _is_clothing_item(item: dict) -> bool:
    """Filter out non-clothing items that the model might detect."""
    item_type = item.get("type", "").lower().strip()
    if item_type in GARMENT_TYPES:
        return True
    for garment in GARMENT_TYPES:
        if garment in item_type or item_type in garment:
            return True
    return False


def _normalize_garment(item: dict) -> dict:
    """Ensure all required fields exist and are normalized."""
    return {
        "type": item.get("type", "unknown").lower().strip(),
        "color": item.get("color", "unknown").lower().strip(),
        "material": item.get("material", "unknown").lower().strip(),
        "pattern": item.get("pattern", "solid").lower().strip(),
        "season": item.get("season", "all").lower().strip(),
        "formality": item.get("formality", "casual").lower().strip(),
        "description": item.get("description", "").strip(),
    }


def _extract_single_garment(crop_bytes: bytes) -> dict | None:
    """Send a single crop to the VLM and extract one garment."""
    llm = model_manager.get_vision_model()
    data_uri = _image_bytes_to_data_uri(crop_bytes)

    response = llm.create_chat_completion(
        messages=[{
            "role": "user",
            "content": [
                {"type": "text", "text": SINGLE_GARMENT_PROMPT},
                {"type": "image_url", "image_url": {"url": data_uri}},
            ],
        }],
        max_tokens=512,
        temperature=0.1,
    )

    raw_text = response["choices"][0]["message"]["content"]
    logger.debug("VLM single-garment response: %s", raw_text[:200])

    items = _parse_json_response(raw_text)
    if not items:
        return None

    item = items[0]
    if not _is_clothing_item(item):
        return None

    return _normalize_garment(item)


def _extract_from_full_image(image_path: str) -> tuple[list[dict], bytes]:
    """Fallback: extract multiple garments from the full image (no YOLO)."""
    llm = model_manager.get_vision_model()
    data_uri, image_bytes = _prepare_image(image_path)

    response = llm.create_chat_completion(
        messages=[{
            "role": "user",
            "content": [
                {"type": "text", "text": MULTI_GARMENT_PROMPT},
                {"type": "image_url", "image_url": {"url": data_uri}},
            ],
        }],
        max_tokens=2048,
        temperature=0.1,
    )

    raw_text = response["choices"][0]["message"]["content"]
    logger.debug("VLM multi-garment response: %s", raw_text[:200])

    items = _parse_json_response(raw_text)
    garments = [_normalize_garment(item) for item in items if _is_clothing_item(item)]

    return garments, image_bytes


def extract_garments(image_path: str) -> list[tuple[dict, bytes]]:
    """Extract garments from an image with individual crops.

    Uses YOLO to detect garment bounding boxes, crops each one, then
    sends each crop to the VLM individually for attribute extraction.

    Falls back to full-image analysis if YOLO detects nothing.

    Returns list of (garment_dict, crop_jpeg_bytes) tuples.
    """
    logger.info("Processing image: %s", image_path)

    crops = detect_and_crop(image_path)

    if crops:
        results = []
        for i, crop_bytes in enumerate(crops):
            logger.info("Analyzing crop %d/%d", i + 1, len(crops))
            garment = _extract_single_garment(crop_bytes)
            if garment:
                results.append((garment, crop_bytes))
        if results:
            logger.info("Extracted %d garments from %d crops", len(results), len(crops))
            return results

    logger.info("Falling back to full-image analysis")
    garments, full_bytes = _extract_from_full_image(image_path)
    return [(g, full_bytes) for g in garments]


def extract_single_from_path(image_path: str) -> list[tuple[dict, bytes]]:
    """Extract a single garment from a photo (no YOLO, direct VLM).

    Use when the user photographs one garment at a time.
    Returns a list with 0 or 1 (garment_dict, image_bytes) tuples.
    """
    logger.info("Single-garment mode: %s", image_path)
    _, image_bytes = _prepare_image(image_path)
    garment = _extract_single_garment(image_bytes)

    if garment:
        return [(garment, image_bytes)]
    return []


def extract_from_crop_bytes(crop_bytes: bytes) -> tuple[dict, bytes] | None:
    """Extract a garment from pre-cropped JPEG bytes (manual bbox).

    Returns (garment_dict, crop_bytes) or None if extraction fails.
    """
    garment = _extract_single_garment(crop_bytes)
    if garment:
        return (garment, crop_bytes)
    return None