#!/usr/bin/env python3
"""
Generate scene frames for microdrama using NanoBanana.

Reads script JSON and generates first/last frames for each scene.
Uses character and object reference sheets for consistency.

Supports both old format (nanobanana_hints) and new format (frames.first/last).
"""

import os
import sys
import json
import base64
import requests
from pathlib import Path
import time

from dotenv import load_dotenv
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(env_path)

BASE_URL = "https://api.laozhang.ai/v1beta"


def call_nanobanana(
    api_key: str,
    prompt: str,
    reference_image: bytes = None,
    reference_mime: str = "image/jpeg",
    aspect_ratio: str = "9:16",
    max_retries: int = 3
) -> tuple:
    """Call NanoBanana to generate an image with retry logic."""
    parts = []

    if reference_image:
        ref_b64 = base64.b64encode(reference_image).decode('utf-8')
        parts.append({"inline_data": {"mime_type": reference_mime, "data": ref_b64}})

    parts.append({"text": prompt})

    payload = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["IMAGE"],
            "imageConfig": {
                "aspectRatio": aspect_ratio,
                "imageSize": "1K"
            }
        }
    }

    for attempt in range(max_retries):
        try:
            response = requests.post(
                f"{BASE_URL}/models/gemini-3-pro-image-preview:generateContent",
                headers={
                    "x-goog-api-key": api_key,
                    "Content-Type": "application/json"
                },
                json=payload,
                timeout=180
            )

            if response.status_code in (502, 503, 504):
                if attempt < max_retries - 1:
                    print(f"    Retry {attempt + 1}/{max_retries} after HTTP {response.status_code}")
                    time.sleep(5)
                    continue

            if response.status_code != 200:
                print(f"    Error: {response.status_code} - {response.text[:200]}")
                return None, None

            result = response.json()

            for candidate in result.get("candidates", []):
                for part in candidate.get("content", {}).get("parts", []):
                    if "inlineData" in part:
                        data = part["inlineData"]
                        return base64.b64decode(data["data"]), data.get("mimeType", "image/png")

            return None, None

        except requests.exceptions.Timeout:
            if attempt < max_retries - 1:
                print(f"    Retry {attempt + 1}/{max_retries} after timeout")
                time.sleep(5)
            else:
                print(f"    Failed after {max_retries} retries (timeout)")
                return None, None

        except Exception as e:
            if attempt < max_retries - 1:
                print(f"    Retry {attempt + 1}/{max_retries} after error: {e}")
                time.sleep(5)
            else:
                print(f"    Failed after {max_retries} retries: {e}")
                return None, None

    return None, None


def build_frame_prompt(
    scene: dict,
    frame_type: str,
    style_context: str = "",
    visual_style: str = "",
    characters_config: dict = None,
    objects_config: dict = None,
    locations_config: dict = None
) -> str:
    """
    Build prompt for frame generation with full context.

    Includes descriptions of all characters, objects, and locations
    used in the scene for visual consistency.
    """
    # Support both old format (nanobanana_hints) and new format (frames)
    if "frames" in scene:
        frame_desc = scene["frames"].get(frame_type, "")
    else:
        hints = scene.get("nanobanana_hints", {})
        frame_desc = hints.get(f"{frame_type}_frame", "")

    # Get caption for this frame (caption_first or caption_last)
    caption_key = f"caption_{frame_type}"
    caption_text = scene.get(caption_key, "")

    # Build visual style context
    style_hints = ""
    if visual_style == "home_video":
        style_hints = "\n- VHS home video aesthetic, nostalgic 90s film grain\n- Warm faded colors, authentic amateur footage feel"
    elif visual_style == "documentary":
        style_hints = "\n- Documentary style, professional but intimate"

    # Gather character descriptions for this scene
    char_descriptions = []
    scene_chars = scene.get("characters", [])
    if characters_config:
        for char_id in scene_chars:
            if char_id in characters_config:
                char_data = characters_config[char_id]
                desc = char_data.get("description", "")
                if desc:
                    char_descriptions.append(f"- {char_id.upper()}: {desc}")

    # Gather object descriptions for this scene
    obj_descriptions = []
    scene_objects = scene.get("objects", [])
    if objects_config:
        for obj_id in scene_objects:
            if obj_id in objects_config:
                obj_data = objects_config[obj_id]
                desc = obj_data.get("description", "")
                if desc:
                    obj_descriptions.append(f"- {obj_id.upper()}: {desc}")

    # Gather location descriptions for this scene
    loc_descriptions = []
    scene_locations = scene.get("locations", [])
    if locations_config:
        for loc_id in scene_locations:
            if loc_id in locations_config:
                loc_data = locations_config[loc_id]
                desc = loc_data.get("description", "")
                if desc:
                    loc_descriptions.append(f"- {loc_id.upper()}: {desc}")

    # Build context sections
    context_sections = []

    if char_descriptions:
        context_sections.append("CHARACTERS IN THIS SCENE (maintain exact appearance from reference):\n" + "\n".join(char_descriptions))

    if obj_descriptions:
        context_sections.append("OBJECTS IN THIS SCENE (maintain exact visual details):\n" + "\n".join(obj_descriptions))

    if loc_descriptions:
        context_sections.append("LOCATION (maintain consistent environment):\n" + "\n".join(loc_descriptions))

    context_text = "\n\n".join(context_sections)
    if context_text:
        context_text = "\n\n" + context_text

    # Build caption instruction if caption exists
    caption_instruction = ""
    if caption_text:
        caption_instruction = f"""

TEXT OVERLAY (render this text on the frame):
- Text: "{caption_text}"
- Position: Bottom of frame, centered
- Style: White text with black outline/shadow for readability
- Font: Clean, readable sans-serif
- Size: Large enough to read on mobile (about 5% of frame height)"""

    prompt = f"""Generate a single photorealistic video frame.

FRAME DESCRIPTION: {frame_desc}{context_text}

VISUAL STYLE:
- Aspect ratio: 9:16 (vertical video)
- Photorealistic, natural lighting
- Authentic, believable scene{style_hints}{style_context}

IMPORTANT:
- Match the reference image style exactly
- Maintain consistent character appearance
- Objects must match their descriptions precisely
- Location details must be consistent{caption_instruction}

Output a crisp, photorealistic frame with the text overlay."""

    return prompt


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("task_dir", help="Task directory path")
    parser.add_argument("--script", default="script_v3.json", help="Script filename")
    parser.add_argument("--scene", "-s", type=int, help="Generate only this scene number")
    parser.add_argument("--frame", "-f", choices=["first", "last", "both"], default="both")
    args = parser.parse_args()

    task_dir = Path(args.task_dir)
    refs_dir = task_dir / "references" / "sheets"

    # Load script (prefer v6, v5, v4, then v3, then specified)
    for version in ["script_v6.json", "script_v5.json", "script_v4.json"]:
        script_path = task_dir / "script" / version
        if script_path.exists():
            break
    else:
        script_path = task_dir / "script" / args.script
        if not script_path.exists():
            script_path = task_dir / "script" / "script.json"

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

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

    visual_style = script.get("visual_style", "")

    # Load ALL references (characters + objects + locations)
    all_refs = {}

    def get_sheet_filename(data):
        """Get sheet filename from data (supports both 'sheet' and 'reference' fields)."""
        if isinstance(data, dict):
            return data.get("sheet") or data.get("reference")
        return data  # Old format: just filename

    # Load character sheets
    characters_config = script.get("characters", {})
    print("\nLoading character references...")
    for char_id, char_data in characters_config.items():
        filename = get_sheet_filename(char_data)
        if filename:
            char_path = refs_dir / filename
            if char_path.exists():
                with open(char_path, "rb") as f:
                    all_refs[char_id] = f.read()
                print(f"  {char_id}: {filename} ({len(all_refs[char_id]) // 1024}KB)")
            else:
                print(f"  {char_id}: {filename} (NOT FOUND)")

    # Load storyteller if present
    storyteller_config = script.get("storyteller", {})
    if storyteller_config:
        filename = get_sheet_filename(storyteller_config)
        if filename:
            st_path = refs_dir / filename
            if st_path.exists():
                with open(st_path, "rb") as f:
                    all_refs["storyteller"] = f.read()
                print(f"  storyteller: {filename} ({len(all_refs['storyteller']) // 1024}KB)")

    # Load object sheets
    objects_config = script.get("objects", {})
    print("\nLoading object references...")
    for obj_id, obj_data in objects_config.items():
        filename = get_sheet_filename(obj_data)
        if filename:
            obj_path = refs_dir / filename
            if obj_path.exists():
                with open(obj_path, "rb") as f:
                    all_refs[obj_id] = f.read()
                print(f"  {obj_id}: {filename} ({len(all_refs[obj_id]) // 1024}KB)")
            else:
                print(f"  {obj_id}: {filename} (NOT FOUND)")

    # Load location sheets
    locations_config = script.get("locations", {})
    print("\nLoading location references...")
    for loc_id, loc_data in locations_config.items():
        filename = get_sheet_filename(loc_data)
        if filename:
            loc_path = refs_dir / filename
            if loc_path.exists():
                with open(loc_path, "rb") as f:
                    all_refs[loc_id] = f.read()
                print(f"  {loc_id}: {filename} ({len(all_refs[loc_id]) // 1024}KB)")
            else:
                print(f"  {loc_id}: {filename} (NOT FOUND)")

    api_key = os.getenv("LAOZHANG_KEY")
    if not api_key:
        print("Error: Set LAOZHANG_KEY")
        sys.exit(1)

    scenes = script.get("scenes", [])
    if args.scene:
        scenes = [s for s in scenes if s["scene_id"] == args.scene]

    print(f"\n{'='*50}")
    print(f"Generating frames for {len(scenes)} scene(s)")
    print(f"Visual style: {visual_style}")
    print(f"{'='*50}\n")

    # Add storyteller to characters_config for reference lookup
    full_characters_config = dict(characters_config)
    if storyteller_config:
        full_characters_config["storyteller"] = storyteller_config

    for scene in scenes:
        scene_id = scene["scene_id"]
        scene_dir = task_dir / "scenes" / f"scene_{scene_id:03d}"
        scene_dir.mkdir(parents=True, exist_ok=True)

        scene_type = scene.get("type", "unknown")
        print(f"=== Scene {scene_id}: {scene_type} ===")

        # Determine PRIMARY reference image to use
        # Priority: character > location > object
        ref_image = None
        ref_mime = "image/jpeg"
        ref_used = None

        # Check characters first
        characters = scene.get("characters", [])
        for char_id in characters:
            if char_id in all_refs:
                ref_image = all_refs[char_id]
                ref_used = f"char:{char_id}"
                break

        # If no character ref, check locations
        if not ref_image:
            locations = scene.get("locations", [])
            for loc_id in locations:
                if loc_id in all_refs:
                    ref_image = all_refs[loc_id]
                    ref_used = f"loc:{loc_id}"
                    break

        # If no location ref, check objects
        if not ref_image:
            objects = scene.get("objects", [])
            for obj_id in objects:
                if obj_id in all_refs:
                    ref_image = all_refs[obj_id]
                    ref_used = f"obj:{obj_id}"
                    break

        print(f"  Primary reference: {ref_used or 'none'}")
        print(f"  Characters: {scene.get('characters', [])}")
        print(f"  Objects: {scene.get('objects', [])}")
        print(f"  Locations: {scene.get('locations', [])}")

        # Style context for sepia scenes
        style_context = ""
        if "sepia" in scene.get("color_grade", "").lower():
            style_context = "\n- SEPIA VINTAGE TONE for 1867 historical period\n- Film grain texture"

        # Generate first frame
        if args.frame in ["first", "both"]:
            print(f"  Generating first frame...")
            prompt = build_frame_prompt(
                scene=scene,
                frame_type="first",
                style_context=style_context,
                visual_style=visual_style,
                characters_config=full_characters_config,
                objects_config=objects_config,
                locations_config=locations_config
            )

            with open(scene_dir / "nb_prompt_first.txt", "w") as f:
                f.write(prompt)

            img_data, img_mime = call_nanobanana(
                api_key, prompt, ref_image, ref_mime, "9:16"
            )

            if img_data:
                ext = "jpg" if "jpeg" in (img_mime or "") else "png"
                frame_path = scene_dir / f"frame_first.{ext}"
                with open(frame_path, "wb") as f:
                    f.write(img_data)
                print(f"    Saved: {frame_path.name} ({len(img_data) // 1024}KB)")
            else:
                print(f"    FAILED")

            time.sleep(2)

        # Generate last frame
        if args.frame in ["last", "both"]:
            print(f"  Generating last frame...")
            prompt = build_frame_prompt(
                scene=scene,
                frame_type="last",
                style_context=style_context,
                visual_style=visual_style,
                characters_config=full_characters_config,
                objects_config=objects_config,
                locations_config=locations_config
            )

            with open(scene_dir / "nb_prompt_last.txt", "w") as f:
                f.write(prompt)

            # For last frame, use first frame as reference for continuity
            frame_ref = ref_image
            first_frame_path = scene_dir / "frame_first.jpg"
            if not first_frame_path.exists():
                first_frame_path = scene_dir / "frame_first.png"

            if first_frame_path.exists() and args.frame == "both":
                with open(first_frame_path, "rb") as f:
                    frame_ref = f.read()
                print(f"    Using first frame for scene continuity")

            img_data, img_mime = call_nanobanana(
                api_key, prompt, frame_ref, ref_mime, "9:16"
            )

            if img_data:
                ext = "jpg" if "jpeg" in (img_mime or "") else "png"
                frame_path = scene_dir / f"frame_last.{ext}"
                with open(frame_path, "wb") as f:
                    f.write(img_data)
                print(f"    Saved: {frame_path.name} ({len(img_data) // 1024}KB)")
            else:
                print(f"    FAILED")

            time.sleep(2)

        print()

    print("Done!")


if __name__ == "__main__":
    main()
