backend/app/product_image_matcher.py

"""
AI-powered product image matching for creatives.
Intelligently selects which product reference images (angles, hand poses, backgrounds)
best match each creative's visual strategy and shooting angle requirements.
"""

import json
import logging

from app.llm import call_llm_vision

log = logging.getLogger("uvicorn.error")


def match_product_images_to_creatives(
    creatives: list[dict],
    product_image_urls: list[str],
    max_images_per_creative: int = 3,
) -> dict[int, list[str]]:
    """
    Use AI to match product reference images to each creative based on visual requirements.
    
    Analyzes each creative's shooting angle, scene requirements, and composition needs,
    then selects the most appropriate product images (e.g., hand on dark bg, flat lay, etc.)
    
    Args:
        creatives: List of creative dicts with scene_prompt, shooting_angle, visual_strategy
        product_image_urls: List of product image URLs (different angles/backgrounds)
        max_images_per_creative: Max number of product images per creative (default 3)
    
    Returns:
        Dict mapping creative_number to list of selected product image URLs
    """
    if not product_image_urls or not creatives:
        return {}
    
    if len(product_image_urls) <= max_images_per_creative:
        # If we have few product images, just use all of them for each creative
        default_selection = product_image_urls[:max_images_per_creative]
        return {
            creative.get("creative_number", i + 1): default_selection
            for i, creative in enumerate(creatives)
        }
    
    log.info("product_matcher: analyzing %d creatives against %d product images", 
             len(creatives), len(product_image_urls))
    
    # Build creative summaries for matching
    creative_summaries = []
    for i, creative in enumerate(creatives[:10]):  # Batch of 10
        vs = creative.get("visual_strategy") or {}
        creative_num = creative.get("creative_number", i + 1)
        archetype = creative.get("archetype", "")
        scene_prompt = (creative.get("scene_prompt") or vs.get("scene_prompt", ""))[:400]
        shooting_angle = vs.get("shooting_angle", "")
        color_world = vs.get("color_world", "")
        
        creative_summaries.append({
            "creative_number": creative_num,
            "archetype": archetype,
            "scene_prompt": scene_prompt,
            "shooting_angle": shooting_angle,
            "color_world": color_world,
        })
    
    prompt = f"""You are an expert product photographer analyzing jewelry product images and ad creative briefs.

TASK: For each creative concept, select the {max_images_per_creative} MOST SUITABLE product reference images.

PRODUCT IMAGES AVAILABLE ({len(product_image_urls)} images total):
Images numbered 1-{len(product_image_urls)} show the same jewelry product from different:
- Angles (overhead, 45-degree, side view, closeup)
- Contexts (on hand with dark bg, on hand with light bg, flat lay, on model)
- Lighting (dramatic, soft, natural, studio)
- Backgrounds (black, white, textured, neutral)

CREATIVES TO MATCH ({len(creative_summaries)} creatives):
{json.dumps(creative_summaries, indent=2)}

MATCHING CRITERIA:
1. **Shooting Angle Match**: If creative needs "overhead" → select overhead product shots
2. **Background/Color Match**: If creative needs "dark moody" → select images with dark backgrounds
3. **Context Match**: If creative needs "lifestyle/hand" → select hand-worn images; if "flat lay" → select flat lay images
4. **Composition**: Match image composition to scene requirements

For EACH creative, select {max_images_per_creative} product image numbers (1-{len(product_image_urls)}) that best match its visual requirements.

Return JSON:
{{
  "matches": {{
    "1": [2, 5, 7],
    "2": [1, 3, 9],
    ...
  }}
}}

IMPORTANT: 
- Match based on VISUAL COMPATIBILITY (angle, lighting, context)
- Prioritize images that align with the creative's shooting angle and scene requirements
- Different creatives may need different product images based on their briefs
"""

    # Build content with product images
    content: list[dict] = [{"type": "text", "text": prompt}]
    for url in product_image_urls:
        content.append({
            "type": "image_url",
            "image_url": {"url": url}
        })
    
    messages = [{"role": "user", "content": content}]
    
    try:
        raw = call_llm_vision(
            messages=messages,
            model="gpt-4o",
            temperature=0.2,  # Very low temperature for consistent product matching
            response_format={"type": "json_object"},
        )
        
        result = json.loads(raw)
        matches = result.get("matches", {})
        
        # Convert image numbers to URLs
        creative_to_images: dict[int, list[str]] = {}
        for creative_num_str, image_numbers in matches.items():
            creative_num = int(creative_num_str)
            selected_urls = []
            for img_num in image_numbers[:max_images_per_creative]:
                idx = img_num - 1  # Convert 1-indexed to 0-indexed
                if 0 <= idx < len(product_image_urls):
                    selected_urls.append(product_image_urls[idx])
            if selected_urls:
                creative_to_images[creative_num] = selected_urls
        
        log.info("product_matcher: matched %d creatives to product images", len(creative_to_images))
        return creative_to_images
        
    except Exception as e:
        log.error("product_matcher: AI matching failed: %s", str(e))
        # Fallback: return all product images for each creative
        fallback = product_image_urls[:max_images_per_creative]
        return {
            creative.get("creative_number", i + 1): fallback
            for i, creative in enumerate(creatives)
        }


def batch_match_product_images(
    creatives: list[dict],
    product_image_urls: list[str],
    batch_size: int = 10,
    max_images_per_creative: int = 3,
) -> dict[int, list[str]]:
    """
    Process large numbers of creatives in batches to match product images.
    
    Args:
        creatives: Full list of creatives
        product_image_urls: List of product image URLs
        batch_size: Number of creatives to process per batch
        max_images_per_creative: Max product images per creative
    
    Returns:
        Dict mapping creative_number to list of selected product image URLs
    """
    if len(product_image_urls) <= max_images_per_creative:
        # If few product images, use all for each creative (no need for AI matching)
        default_selection = product_image_urls[:max_images_per_creative]
        return {
            creative.get("creative_number", i + 1): default_selection
            for i, creative in enumerate(creatives)
        }
    
    all_matches = {}
    
    for i in range(0, len(creatives), batch_size):
        batch = creatives[i:i + batch_size]
        log.info("product_matcher: processing batch %d-%d of %d creatives",
                 i + 1, min(i + batch_size, len(creatives)), len(creatives))
        
        batch_matches = match_product_images_to_creatives(
            batch,
            product_image_urls,
            max_images_per_creative=max_images_per_creative,
        )
        all_matches.update(batch_matches)
    
    return all_matches