#!/usr/bin/env python3
"""
Auto-select best product image and generate prompts using Claude Code

Usage:
    python auto_select_image.py <WB_URL>
    python auto_select_image.py --dir <path_to_images> --name "Product Name"
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path

import requests


def download_images(image_urls: list, output_dir: Path, limit: int = 5) -> tuple:
    """Download images from URLs to local directory

    Returns:
        tuple: (list of local paths, dict mapping local path -> original URL)
    """
    downloaded = []
    url_map = {}  # local_path -> original_url
    for i, url in enumerate(image_urls[:limit]):
        try:
            resp = requests.get(url, timeout=10)
            if resp.status_code == 200:
                ext = url.split('.')[-1].split('?')[0]
                filename = f"image_{i+1}.{ext}"
                filepath = output_dir / filename
                filepath.write_bytes(resp.content)
                downloaded.append(str(filepath))
                url_map[str(filepath)] = url
                print(f"  Downloaded: {filename}")
        except Exception as e:
            print(f"  Failed to download {url}: {e}")
    return downloaded, url_map


def run_claude_code(image_paths: list, product_name: str) -> dict:
    """Run Claude Code to analyze images and generate prompts"""

    # Build image list for prompt
    image_list = "\n".join([f"- {path}" for path in image_paths])

    # Build prompt - Claude will use Read tool to view images
    prompt = f'''You are a creative director making product videos that CREATE DESIRE through storytelling.

Product: {product_name}

Image files to analyze:
{image_list}

READ each image, then create a TWO-ACT story:

## THE RULE: Every video must show the product being USED with a POSITIVE OUTCOME.

**ACT 1** (first 3 seconds): The product - beautiful, desirable, waiting
**ACT 2** (last 5 seconds): The product IN USE - a person enjoying it, a positive moment, the BENEFIT

Examples of good Act 2 ideas:
- Suitcase → Person rolling it through airport, arriving at destination, hand grabbing handle
- Jewelry → Woman smiling in mirror, receiving compliment, confident posture
- Cookware → Family gathered around table with FINISHED meal, person holding pan proudly, gift-giving/unboxing moment (AVOID: cooking action, cracking eggs, sliding food, pouring)
- Clothing → Person walking confidently, turning heads, feeling great
- Tech → Person relaxing with music, productive moment, satisfied expression

## VEO3 LIMITATIONS - AVOID THESE ACTIONS:
- Sliding/pouring/cracking (eggs sliding off pan, pouring liquid)
- Complex hand manipulation of small objects
- Rapid scene transitions (prefer single continuous scene or simple fade)
- Objects changing shape or state

## NANOBANANA PHYSICS:
- Products must REST FIRMLY ON surfaces with proper shadows
- Say "resting firmly on the wooden counter" not just "on counter"
- Ensure visible contact point with surface

## 1. SELECT IMAGES:

**detail_image**: Where product is LARGEST (NanoBanana needs pixels!)
**context_image**: Any with lifestyle/human element, or same if none

## 2. CREATIVE CONCEPT:
What positive emotion or outcome does this product enable? Joy? Confidence? Adventure? Comfort? Success?

## 3. NANOBANANA PROMPTS:

**frame_act1**: The product hero shot. Beautiful, inviting, desirable.
Describe: setting, lighting, mood, atmosphere. Make it aspirational.

**frame_act2**: The LIFESTYLE payoff. Product being USED by a PERSON with POSITIVE emotion.
Describe: who is using it, where, what are they feeling, what's the positive outcome.
This is NOT just the product again - it's someone ENJOYING the product!

End each with: "16:9 aspect ratio, photorealistic, no text or watermarks"

## 4. VEO3 PROMPT:
Write a TWO-ACT scene (not just camera movement!):
- ACT 1: Establish the beautiful product
- TRANSITION: How do we move to the lifestyle scene? (fade, match cut, hand enters, etc.)
- ACT 2: Show the product IN USE with positive emotion

The viewer should WANT this product after watching.
End with --no text, watermarks, morphing, distortion, artifacts

Output ONLY valid JSON:
{{
  "detail_image": "filename",
  "context_image": "filename",
  "positive_outcome": "What benefit/emotion does this product provide?",
  "frame_act1": "NanoBanana prompt - product hero shot",
  "frame_act2": "NanoBanana prompt - lifestyle with person USING product",
  "veo3_prompt": "Two-act scene with transition and positive payoff"
}}'''

    # Build claude command (--allowedTools Read enables image viewing)
    cmd = ["claude", "-p", "--allowedTools", "Read"]

    print(f"\nRunning Claude Code to analyze {len(image_paths)} images...")

    try:
        result = subprocess.run(
            cmd,
            input=prompt,  # Pass prompt via stdin
            capture_output=True,
            text=True,
            timeout=180
        )

        if result.returncode != 0:
            print(f"Error: {result.stderr}")
            return None

        # Parse JSON from output
        output = result.stdout.strip()

        # Try to extract JSON from output
        try:
            # Find JSON in output (might have extra text)
            start = output.find('{')
            end = output.rfind('}') + 1
            if start >= 0 and end > start:
                json_str = output[start:end]
                return json.loads(json_str)
        except json.JSONDecodeError:
            print(f"Could not parse JSON from output:\n{output}")
            return {"raw_output": output}

    except subprocess.TimeoutExpired:
        print("Timeout waiting for Claude Code")
        return None
    except FileNotFoundError:
        print("Claude Code not found. Make sure 'claude' is in PATH")
        return None


def main():
    parser = argparse.ArgumentParser(description="Auto-select product image with Claude Code")
    parser.add_argument("url", nargs="?", help="Wildberries product URL")
    parser.add_argument("--dir", help="Directory with downloaded images")
    parser.add_argument("--name", default="Product", help="Product name")
    parser.add_argument("--limit", type=int, default=7, help="Max images to analyze (default: 7)")
    parser.add_argument("--keep", action="store_true", help="Keep downloaded images")

    args = parser.parse_args()

    if not args.url and not args.dir:
        parser.print_help()
        sys.exit(1)

    image_paths = []
    url_map = {}  # Maps local path -> original URL
    temp_dir = None
    product_name = args.name

    try:
        if args.dir:
            # Use existing directory
            dir_path = Path(args.dir)
            image_paths = sorted([
                str(p) for p in dir_path.glob("*")
                if p.suffix.lower() in ['.jpg', '.jpeg', '.png', '.webp']
            ])[:args.limit]

        elif args.url:
            # Fetch from WB and download images
            from wb_fetcher import fetch_wb_product_playwright

            print(f"Fetching product data from: {args.url}")
            product_data = fetch_wb_product_playwright(args.url, headless=True)

            if not product_data:
                print("Failed to fetch product data")
                sys.exit(1)

            product_name = product_data.get("name", "Product")
            image_urls = product_data.get("images", [])

            if not image_urls:
                print("No images found")
                sys.exit(1)

            print(f"\nProduct: {product_name}")
            print(f"Found {len(image_urls)} images, downloading first {args.limit}...")

            # Create temp directory for images
            temp_dir = tempfile.mkdtemp(prefix="wb_images_")
            image_paths, url_map = download_images(image_urls, Path(temp_dir), args.limit)

        if not image_paths:
            print("No images to analyze")
            sys.exit(1)

        # Run Claude Code
        result = run_claude_code(image_paths, product_name)

        if result:
            # Add full paths and original URLs for selected images
            for key in ["detail_image", "context_image"]:
                selected = result.get(key, "")
                for img_path in image_paths:
                    if selected in img_path:
                        result[f"{key}_path"] = img_path
                        if img_path in url_map:
                            result[f"{key}_url"] = url_map[img_path]
                        break

            print("\n" + "="*50)
            print("RESULT:")
            print("="*50)
            print(json.dumps(result, indent=2, ensure_ascii=False))

            # Save result
            output_file = Path("selection_result.json")
            output_file.write_text(json.dumps(result, indent=2, ensure_ascii=False))
            print(f"\nSaved to: {output_file}")

    finally:
        # Cleanup temp directory
        if temp_dir and not args.keep:
            import shutil
            shutil.rmtree(temp_dir, ignore_errors=True)
            print(f"\nCleaned up temp directory")
        elif temp_dir:
            print(f"\nImages kept in: {temp_dir}")


if __name__ == "__main__":
    main()
