routers/flow.py

import asyncio
import logging
from fastapi import APIRouter, Request, HTTPException
from pydantic import BaseModel

logger = logging.getLogger(__name__)

router = APIRouter()


# ── Request models ────────────────────────────────────────────────────────────

class AnalyzeFlowRequest(BaseModel):
    imageBase64: str
    mimeType: str


class SearchFlowRequest(BaseModel):
    imageBase64: str
    mimeType: str


class RecommendFlowRequest(BaseModel):
    productId: str


# ── Helpers ───────────────────────────────────────────────────────────────────

def _build_filter(fashion_attrs: dict) -> dict | None:
    clauses = []
    if fashion_attrs.get("graphical_appearance"):
        clauses.append({"graphical_appearance": {"$eq": fashion_attrs["graphical_appearance"]}})
    if fashion_attrs.get("product_group"):
        clauses.append({"product_group": {"$eq": fashion_attrs["product_group"]}})
    if fashion_attrs.get("index_group"):
        clauses.append({"index_group": {"$eq": fashion_attrs["index_group"]}})
    return {"$and": clauses} if clauses else None


# ── Step 1: Analyze ──────────────────────────────────────────────────────────

@router.post("/analyze")
async def analyze(request: Request, dto: AnalyzeFlowRequest):
    """
    Step 1 — Analyze uploaded image.
    Claude classifies first (fashion / food / unknown).
    If fashion, YOLO detects and crops individual items.
    """
    yolo_service = request.app.state.yolo_service
    vision_service = request.app.state.vision_service

    logger.info("Step 1 — analyze: asking Claude to classify image type")
    loop = asyncio.get_event_loop()

    # Claude classifies first — YOLO is fashion-only and will hallucinate on non-fashion images
    claude_result = await loop.run_in_executor(
        None, vision_service.analyze_image, dto.imageBase64, dto.mimeType
    )
    logger.info(f"Claude vision result type: {claude_result.get('type')}")

    if claude_result.get("type") == "food":
        logger.info(f"Food detected: {claude_result.get('foodAnalysis', {}).get('dishName')}")
        return {"type": "food", "foodAnalysis": claude_result.get("foodAnalysis")}

    if claude_result.get("type") != "fashion":
        logger.warning("Image could not be classified as fashion or food — returning unknown")
        return {"type": "unknown"}

    # Confirmed fashion — run YOLO to detect and crop individual items
    logger.info("Fashion confirmed — running YOLO detection")
    yolo_result = await loop.run_in_executor(
        None, yolo_service.detect_and_crop, dto.imageBase64, dto.mimeType
    )
    detected_items = yolo_result.get("detectedItems", [])
    logger.info(f"YOLO detected {len(detected_items)} item(s)")

    real_items = [i for i in detected_items if i.get("label") != "unknown"]

    if real_items:
        # Deduplicate by label — keep highest confidence per label
        by_label: dict = {}
        for item in real_items:
            label = item["label"]
            if label not in by_label or item["confidence"] > by_label[label]["confidence"]:
                by_label[label] = item
        deduped = list(by_label.values())
        logger.info(f"Fashion items ({len(real_items)} → {len(deduped)} after dedup)")
        return {
            "type": "fashion",
            "items": [
                {
                    "label": item["label"],
                    "confidence": item["confidence"],
                    "boundingBox": item["boundingBox"],
                    "croppedImageBase64": item["croppedImageBase64"],
                    "mimeType": item["mimeType"],
                }
                for item in deduped
            ],
        }

    # YOLO found nothing despite fashion classification — return full image as single item
    logger.warning("YOLO found no items in confirmed fashion image — returning full image")
    fashion_attrs = claude_result.get("fashionAttributes", {})
    return {
        "type": "fashion",
        "items": [{
            "label": fashion_attrs.get("category", "unknown"),
            "confidence": 1.0,
            "boundingBox": {
                "x": 0,
                "y": 0,
                "width": yolo_result.get("imageWidth", 0),
                "height": yolo_result.get("imageHeight", 0),
            },
            "croppedImageBase64": dto.imageBase64,
            "mimeType": dto.mimeType,
        }],
    }


# ── Step 2: Search ────────────────────────────────────────────────────────────

@router.post("/search")
async def search(request: Request, dto: SearchFlowRequest):
    """
    Step 2 — Search by selected crop.
    Generates CLIP image vector + Claude fashion attributes in parallel,
    then queries Pinecone (top 6).
    """
    clip_service = request.app.state.clip_service
    vision_service = request.app.state.vision_service
    pinecone_service = request.app.state.pinecone_service

    logger.info("Step 2 — search: generating CLIP vector and analyzing fashion attributes in parallel")
    loop = asyncio.get_event_loop()

    vector, fashion_attrs = await asyncio.gather(
        loop.run_in_executor(None, clip_service.generate_image_vector, dto.imageBase64),
        loop.run_in_executor(None, vision_service.analyze_fashion_for_search, dto.imageBase64, dto.mimeType),
    )

    logger.info(f"Fashion attributes from Claude: {fashion_attrs}")

    pinecone_filter = _build_filter(fashion_attrs)
    logger.info(f"Pinecone filter: {pinecone_filter}")

    results = await loop.run_in_executor(
        None, lambda: pinecone_service.query(vector, 6, pinecone_filter)
    )

    logger.info(f"Results: {len(results)} from Pinecone")

    return {
        "fashionAttributes": fashion_attrs,
        "results": [
            {"id": r["id"], "score": r["score"], "metadata": r["metadata"]}
            for r in results
        ],
    }


# ── Step 3: Recommend ─────────────────────────────────────────────────────────

@router.post("/recommend")
async def recommend(request: Request, dto: RecommendFlowRequest):
    """
    Step 3 — Get outfit recommendations for a product.
    Fetches product from Pinecone, generates outfit ideas via Claude,
    then matches each outfit item to a real product via CLIP text vector + Pinecone.
    """
    clip_service = request.app.state.clip_service
    vision_service = request.app.state.vision_service
    pinecone_service = request.app.state.pinecone_service

    logger.info(f"Step 3 — recommend: fetching product \"{dto.productId}\" from Pinecone")
    loop = asyncio.get_event_loop()

    product = await loop.run_in_executor(None, pinecone_service.fetch_by_id, dto.productId)
    if not product:
        logger.error(f"Product \"{dto.productId}\" not found in Pinecone")
        raise HTTPException(status_code=404, detail=f"Product \"{dto.productId}\" not found in Pinecone")

    meta = product.get("metadata", {})
    logger.info(f"Product found: {meta.get('name')}")

    description_parts = [
        meta.get("name"),
        f"in {meta['colour']}" if meta.get("colour") else None,
        f"— {meta['description']}" if meta.get("description") else None,
    ]
    description = " ".join(p for p in description_parts if p)

    logger.info(f"Requesting outfit ideas from Claude for: \"{description}\"")
    outfit_result = await loop.run_in_executor(
        None, vision_service.get_outfit_ideas_from_description, description
    )
    outfit = outfit_result.get("outfit", {})
    items = outfit.get("items", [])
    logger.info(f"Claude suggested outfit: \"{outfit.get('title')}\" with {len(items)} item(s)")

    # Match each outfit item to a real product via CLIP text vector + Pinecone top 1
    logger.info("Matching each outfit item to a real product via CLIP text vector + Pinecone")

    async def match_item(item: dict) -> dict:
        query = f"{item.get('color', '')} {item.get('item', '')}".strip()
        logger.info(f"  Matching outfit item: \"{query}\"")

        item_filter = _build_filter(item)
        logger.info(
            f"  Filters: graphical_appearance={item.get('graphical_appearance', 'none')}, "
            f"product_group={item.get('product_group', 'none')}, "
            f"index_group={item.get('index_group', 'none')}"
        )

        text_vector = await loop.run_in_executor(None, clip_service.generate_text_vector, query)
        matches = await loop.run_in_executor(
            None, lambda: pinecone_service.query(text_vector, 1, item_filter)
        )
        matched = matches[0] if matches else None
        logger.info(
            f"  → Matched: {matched['metadata'].get('name') if matched else 'none'} "
            f"(score: {matched['score'] if matched else '-'})"
        )

        return {
            **item,
            "matchedProduct": (
                {"id": matched["id"], "score": matched["score"], "metadata": matched["metadata"]}
                if matched else None
            ),
        }

    enriched_items = await asyncio.gather(*[match_item(item) for item in items])

    logger.info("Step 3 — recommend complete")
    return {
        "success": True,
        "detectedItem": outfit_result.get("detectedItem"),
        "claudeRawResponse": outfit_result,
        "outfit": {
            **outfit,
            "items": list(enriched_items),
        },
    }