#!/usr/bin/env python3
"""
Download references and generate sheets from existing script.json.

Reads script.json and:
1. Downloads reference images via Yandex for all entities
2. Uses Claude to select best references
3. Generates character/object/location sheets with NanoBanana

Usage:
    python download_references.py /path/to/task_dir
    python download_references.py /path/to/task_dir --skip-download  # Only generate sheets
"""

import os
import sys
import json
import argparse
import base64
import time
from pathlib import Path
from typing import Dict, List, Optional

import requests
from dotenv import load_dotenv

# Load .env
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(env_path)

# Import from parent scripts
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
from laozhang_proxy import call_nanobanana_laozhang

SCRIPTS_DIR = Path(__file__).parent
LAOZHANG_KEY = os.getenv("LAOZHANG_KEY")
LAOZHANG_CLAUDE_URL = "https://api.laozhang.ai/v1/chat/completions"


def call_claude(prompt: str, images: List[Dict] = None, model: str = "claude-sonnet-4-20250514") -> str:
    """Call Claude via LaoZhang API."""
    messages = [{"role": "user", "content": []}]

    if images:
        for img in images:
            messages[0]["content"].append({
                "type": "image_url",
                "image_url": {"url": f"data:{img['mime_type']};base64,{img['base64']}"}
            })

    messages[0]["content"].append({"type": "text", "text": prompt})

    response = requests.post(
        LAOZHANG_CLAUDE_URL,
        headers={
            "Authorization": f"Bearer {LAOZHANG_KEY}",
            "Content-Type": "application/json"
        },
        json={"model": model, "messages": messages, "max_tokens": 2000},
        timeout=120
    )

    if response.status_code != 200:
        print(f"Claude API error: {response.status_code}")
        return ""

    return response.json()["choices"][0]["message"]["content"]


def search_yandex(query: str, output_dir: Path, prefix: str, count: int = 5) -> List[Path]:
    """Search Yandex and download images."""
    import subprocess

    print(f"  Searching: {query}")

    cmd = [
        sys.executable,
        str(SCRIPTS_DIR / "yandex_image_search.py"),
        query, "-o", str(output_dir), "-n", str(count),
        "--size", "large", "--download", "-p", prefix
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
        if result.returncode != 0:
            print(f"    Warning: Search may have partial results")
    except Exception as e:
        print(f"    Error: {e}")
        return []

    # Find downloaded files
    downloaded = list(output_dir.glob(f"{prefix}_*.jpg")) + \
                 list(output_dir.glob(f"{prefix}_*.webp")) + \
                 list(output_dir.glob(f"{prefix}_*.png"))
    print(f"    Downloaded {len(downloaded)} images")
    return sorted(downloaded)


def select_best_reference(images: List[Path], description: str) -> Optional[Path]:
    """Use Claude to select best reference image."""
    if not images:
        return None
    if len(images) == 1:
        return images[0]

    print(f"    Selecting best from {len(images)} images...")

    # Prepare images for Claude
    image_data = []
    for img_path in images[:5]:
        with open(img_path, "rb") as f:
            img_bytes = f.read()
        img_b64 = base64.b64encode(img_bytes).decode()
        mime = "image/jpeg"
        if img_path.suffix.lower() == ".png":
            mime = "image/png"
        elif img_path.suffix.lower() == ".webp":
            mime = "image/webp"
        image_data.append({"base64": img_b64, "mime_type": mime})

    prompt = f"""Select the best reference image for: {description}

Consider: quality, match to description, suitability for video generation.
Respond with ONLY the number (1, 2, 3, etc.) of the best image."""

    try:
        response = call_claude(prompt, images=image_data)
        for word in response.split():
            if word.isdigit():
                idx = int(word) - 1
                if 0 <= idx < len(images):
                    print(f"    Selected: {images[idx].name}")
                    return images[idx]
    except Exception as e:
        print(f"    Selection error: {e}")

    return images[0]


def download_all_references(script_data: Dict, task_dir: Path) -> Dict[str, Path]:
    """Download and select references for all entities in script."""
    search_dir = task_dir / "references" / "search"
    search_dir.mkdir(parents=True, exist_ok=True)

    selected_refs = {}

    # Storyteller
    storyteller = script_data.get("storyteller", {})
    if storyteller.get("reference_query"):
        print("\n[STORYTELLER]")
        images = search_yandex(storyteller["reference_query"], search_dir, "storyteller")
        best = select_best_reference(images, storyteller.get("description", ""))
        if best:
            selected_refs["storyteller"] = best

    # Characters
    for char_id, char_info in script_data.get("characters", {}).items():
        if char_info.get("reference_query"):
            print(f"\n[CHARACTER: {char_id}]")
            images = search_yandex(char_info["reference_query"], search_dir, char_id)
            best = select_best_reference(images, char_info.get("description", ""))
            if best:
                selected_refs[char_id] = best

    # Objects
    for obj_id, obj_info in script_data.get("objects", {}).items():
        if obj_info.get("reference_query"):
            print(f"\n[OBJECT: {obj_id}]")
            images = search_yandex(obj_info["reference_query"], search_dir, obj_id)
            best = select_best_reference(images, obj_info.get("description", ""))
            if best:
                selected_refs[obj_id] = best

    # Locations
    for loc_id, loc_info in script_data.get("locations", {}).items():
        if loc_info.get("reference_query"):
            print(f"\n[LOCATION: {loc_id}]")
            images = search_yandex(loc_info["reference_query"], search_dir, loc_id)
            best = select_best_reference(images, loc_info.get("description", ""))
            if best:
                selected_refs[loc_id] = best

    # Save mapping
    refs_path = task_dir / "references" / "selected_refs.json"
    with open(refs_path, "w") as f:
        json.dump({k: str(v) for k, v in selected_refs.items()}, f, indent=2)

    print(f"\nSelected {len(selected_refs)} reference images")
    return selected_refs


def generate_sheet(
    entity_id: str,
    entity_info: Dict,
    ref_path: Optional[Path],
    sheets_dir: Path,
    prompts_dir: Path,
    entity_type: str = "character"
) -> Optional[Path]:
    """Generate a single reference sheet."""
    sheet_path = sheets_dir / f"sheet_{entity_id}.jpg"
    prompt_path = prompts_dir / f"sheet_{entity_id}_prompt.txt"

    if sheet_path.exists():
        print(f"  Exists: {sheet_path.name}")
        return sheet_path

    description = entity_info.get("description", "")
    is_location = entity_type == "location"

    # Load reference image if available
    ref_data = None
    ref_mime = "image/jpeg"
    if ref_path and ref_path.exists():
        with open(ref_path, "rb") as f:
            ref_data = f.read()
        if ref_path.suffix.lower() == ".png":
            ref_mime = "image/png"
        elif ref_path.suffix.lower() == ".webp":
            ref_mime = "image/webp"

    # Build prompt based on entity type
    if is_location:
        prompt = f"""Create a reference sheet showing this location from multiple angles:

LOCATION: {description}

Layout:
- Main view (large, top): Wide establishing shot of the location
- Three smaller views below: Different angles/details of the same location

CRITICAL: This is a LOCATION reference - NO PEOPLE in the image!
Show only the environment, architecture, furniture, lighting, and atmosphere.
Use a clean, neutral presentation style suitable for video production reference.
Maintain consistent lighting and color palette across all views."""
    else:
        # Character or object
        has_ref = ref_data is not None
        if has_ref:
            prompt = f"""Create a professional character reference sheet based on the uploaded reference image.

CHARACTER: {description}

Layout - Two horizontal rows:
Top row: Four full-body views (front, left profile, right profile, back)
Bottom row: Three close-up portraits (front, left profile, right profile)

Use clean neutral background. Match the reference image style exactly.
Maintain perfect identity consistency across all panels.
Use relaxed A-pose for standing views. Sharp details, even lighting."""
        else:
            prompt = f"""Create a professional character reference sheet:

CHARACTER: {description}

Layout - Two horizontal rows:
Top row: Four full-body views (front, left profile, right profile, back)
Bottom row: Three close-up portraits (front, left profile, right profile)

Use clean neutral background, photorealistic style.
Maintain perfect identity consistency across all panels.
Use relaxed A-pose for standing views. Sharp details, even lighting."""

    # Save prompt
    with open(prompt_path, "w") as f:
        f.write(prompt)
        if ref_path:
            f.write(f"\n\n# Reference: {ref_path}")

    # Generate sheet
    print(f"  Generating {entity_type} sheet: {entity_id}...")
    try:
        sheet_bytes, _ = call_nanobanana_laozhang(
            prompt=prompt,
            image_data=ref_data,
            mime_type=ref_mime,
            api_key=LAOZHANG_KEY,
            aspect_ratio="16:9"
        )

        if sheet_bytes:
            with open(sheet_path, "wb") as f:
                f.write(sheet_bytes)
            print(f"    Saved: {sheet_path.name}")
            return sheet_path
        else:
            print(f"    Failed to generate")
            return ref_path
    except Exception as e:
        print(f"    Error: {e}")
        return ref_path


def generate_all_sheets(script_data: Dict, selected_refs: Dict[str, Path], task_dir: Path) -> Dict[str, Path]:
    """Generate sheets for all entities."""
    sheets_dir = task_dir / "references" / "sheets"
    prompts_dir = task_dir / "references" / "prompts"
    sheets_dir.mkdir(parents=True, exist_ok=True)
    prompts_dir.mkdir(parents=True, exist_ok=True)

    sheets = {}

    # Storyteller
    storyteller = script_data.get("storyteller", {})
    if storyteller:
        print("\n[SHEET: storyteller]")
        result = generate_sheet(
            "storyteller", storyteller,
            selected_refs.get("storyteller"),
            sheets_dir, prompts_dir, "character"
        )
        if result:
            sheets["storyteller"] = result

    # Characters
    for char_id, char_info in script_data.get("characters", {}).items():
        print(f"\n[SHEET: {char_id}]")
        result = generate_sheet(
            char_id, char_info,
            selected_refs.get(char_id),
            sheets_dir, prompts_dir, "character"
        )
        if result:
            sheets[char_id] = result
        time.sleep(2)  # Rate limit

    # Objects
    for obj_id, obj_info in script_data.get("objects", {}).items():
        print(f"\n[SHEET: {obj_id}]")
        result = generate_sheet(
            obj_id, obj_info,
            selected_refs.get(obj_id),
            sheets_dir, prompts_dir, "object"
        )
        if result:
            sheets[obj_id] = result
        time.sleep(2)

    # Locations
    for loc_id, loc_info in script_data.get("locations", {}).items():
        print(f"\n[SHEET: {loc_id}]")
        result = generate_sheet(
            loc_id, loc_info,
            selected_refs.get(loc_id),
            sheets_dir, prompts_dir, "location"
        )
        if result:
            sheets[loc_id] = result
        time.sleep(2)

    # Save mapping
    with open(sheets_dir / "sheets.json", "w") as f:
        json.dump({k: str(v) for k, v in sheets.items()}, f, indent=2)

    # Update script with sheet paths
    update_script_with_sheets(script_data, sheets, task_dir)

    print(f"\nGenerated {len(sheets)} reference sheets")
    return sheets


def update_script_with_sheets(script_data: Dict, sheets: Dict[str, Path], task_dir: Path):
    """Update script.json with sheet file paths."""
    # Update storyteller
    if "storyteller" in sheets and "storyteller" in script_data:
        script_data["storyteller"]["sheet"] = sheets["storyteller"].name

    # Update characters
    for char_id in script_data.get("characters", {}):
        if char_id in sheets:
            script_data["characters"][char_id]["sheet"] = sheets[char_id].name

    # Update objects
    for obj_id in script_data.get("objects", {}):
        if obj_id in sheets:
            script_data["objects"][obj_id]["sheet"] = sheets[obj_id].name

    # Update locations
    for loc_id in script_data.get("locations", {}):
        if loc_id in sheets:
            script_data["locations"][loc_id]["sheet"] = sheets[loc_id].name

    # Save updated script
    script_path = task_dir / "script" / "script.json"
    with open(script_path, "w") as f:
        json.dump(script_data, f, ensure_ascii=False, indent=2)

    print(f"\nUpdated script with sheet paths: {script_path}")


def main():
    parser = argparse.ArgumentParser(description="Download references and generate sheets")
    parser.add_argument("task_dir", help="Task directory with script.json")
    parser.add_argument("--skip-download", action="store_true", help="Skip Yandex download, only generate sheets")
    parser.add_argument("--skip-sheets", action="store_true", help="Only download, don't generate sheets")

    args = parser.parse_args()
    task_dir = Path(args.task_dir)

    # Load script
    script_path = task_dir / "script" / "script.json"
    if not script_path.exists():
        # Try versioned scripts
        for v in ["script_v6.json", "script_v5.json"]:
            alt_path = task_dir / "script" / v
            if alt_path.exists():
                script_path = alt_path
                break

    if not script_path.exists():
        print(f"Error: No script found in {task_dir / 'script'}")
        sys.exit(1)

    print(f"Loading: {script_path}")
    with open(script_path) as f:
        script_data = json.load(f)

    # Download references
    if args.skip_download:
        # Load existing selections
        refs_path = task_dir / "references" / "selected_refs.json"
        if refs_path.exists():
            with open(refs_path) as f:
                selected_refs = {k: Path(v) for k, v in json.load(f).items()}
        else:
            selected_refs = {}
        print(f"Loaded {len(selected_refs)} existing references")
    else:
        print("\n" + "=" * 60)
        print("DOWNLOADING REFERENCES")
        print("=" * 60)
        selected_refs = download_all_references(script_data, task_dir)

    # Generate sheets
    if not args.skip_sheets:
        print("\n" + "=" * 60)
        print("GENERATING SHEETS")
        print("=" * 60)
        generate_all_sheets(script_data, selected_refs, task_dir)

    print("\n" + "=" * 60)
    print("COMPLETE!")
    print("=" * 60)


if __name__ == "__main__":
    main()
