#!/usr/bin/env python3
"""
Automated frame generation for microdrama pipeline.

Reads script JSON and generates all frames:
- Ken Burns scenes: single frame_first.jpg
- Veo3 scenes: frame_first.jpg and frame_last.jpg

Key features:
- Combines ALL relevant references (characters + objects + locations)
- Uses frame_first as reference for frame_last (consistency)
- Consistent caption styling across all frames
"""

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

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

# Import microdrama styles
from microdrama_styles import get_frame_style_hints, get_style_config

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


def call_nanobanana(
    api_key: str,
    prompt: str,
    images: List[Tuple[bytes, str]] = None,  # List of (image_data, mime_type)
    aspect_ratio: str = "9:16",
    max_retries: int = 3
) -> tuple:
    """Call NanoBanana Pro with multiple reference images."""
    parts = []

    # Add all reference images
    if images:
        for image_data, mime_type in images:
            image_base64 = base64.b64encode(image_data).decode('utf-8')
            parts.append({"inline_data": {"mime_type": mime_type, "data": image_base64}})

    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"]
                        image_bytes = base64.b64decode(data["data"])
                        return image_bytes, data.get("mimeType", "image/png")

            print("    No image data in response")
            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 load_image(path: Path) -> Tuple[bytes, str]:
    """Load image and return (data, mime_type)."""
    with open(path, "rb") as f:
        data = f.read()

    mime = "image/jpeg"
    if path.suffix.lower() == ".png":
        mime = "image/png"
    elif path.suffix.lower() == ".webp":
        mime = "image/webp"

    return data, mime


def find_sheet(entity_id: str, entity_info: Dict, sheets_dir: Path) -> Optional[Path]:
    """Find sheet file for an entity."""
    # Check 'sheet' field first (new format)
    sheet_name = entity_info.get("sheet")
    if sheet_name:
        sheet_path = sheets_dir / sheet_name
        if sheet_path.exists():
            return sheet_path

    # Try sheet_entityid.jpg format
    for ext in ["jpg", "jpeg", "png", "webp"]:
        sheet_path = sheets_dir / f"sheet_{entity_id}.{ext}"
        if sheet_path.exists():
            return sheet_path

    # Check 'reference' field (old format)
    ref_name = entity_info.get("reference")
    if ref_name:
        ref_path = sheets_dir / ref_name
        if ref_path.exists():
            return ref_path

    return None


def gather_scene_references(
    scene: Dict,
    script: Dict,
    sheets_dir: Path
) -> List[Tuple[str, Path]]:
    """
    Gather ALL relevant reference sheets for a scene.

    Returns list of (label, path) tuples for:
    - Characters in the scene
    - Objects in the scene
    - Locations in the scene
    """
    refs = []

    # Characters
    for char_id in scene.get("characters", []):
        if char_id == "storyteller":
            char_info = script.get("storyteller", {})
        else:
            char_info = script.get("characters", {}).get(char_id, {})

        if char_info:
            sheet_path = find_sheet(char_id, char_info, sheets_dir)
            if sheet_path:
                refs.append((f"character:{char_id}", sheet_path))

    # Objects
    for obj_id in scene.get("objects", []):
        obj_info = script.get("objects", {}).get(obj_id, {})
        if obj_info:
            sheet_path = find_sheet(obj_id, obj_info, sheets_dir)
            if sheet_path:
                refs.append((f"object:{obj_id}", sheet_path))

    # Locations
    for loc_id in scene.get("locations", []):
        loc_info = script.get("locations", {}).get(loc_id, {})
        if loc_info:
            sheet_path = find_sheet(loc_id, loc_info, sheets_dir)
            if sheet_path:
                refs.append((f"location:{loc_id}", sheet_path))

    return refs


def get_caption_style(script: Dict) -> Dict:
    """
    Determine consistent caption style based on visual_style.
    Returns dict with caption styling instructions.
    """
    visual_style = script.get("visual_style", "documentary")

    # Define caption styles that match visual styles
    styles = {
        "home_video": {
            "position": "bottom center",
            "background": "semi-transparent black bar",
            "font": "clean sans-serif (like Arial)",
            "color": "white",
            "size": "medium, readable on mobile",
            "effect": "slight VHS noise/grain overlay",
            "description": "Simple white text on dark bar, VHS-era style"
        },
        "documentary": {
            "position": "bottom center",
            "background": "subtle dark gradient",
            "font": "clean modern sans-serif",
            "color": "white with subtle shadow",
            "size": "medium",
            "effect": "none",
            "description": "Professional documentary lower-third style"
        },
        "news_broadcast": {
            "position": "bottom, news ticker style",
            "background": "solid colored bar (red/blue)",
            "font": "bold sans-serif",
            "color": "white",
            "size": "large, bold",
            "effect": "news ticker animation hint",
            "description": "News channel lower-third with colored background"
        },
        "melodrama": {
            "position": "center or bottom",
            "background": "none or subtle shadow",
            "font": "elegant serif or script",
            "color": "white with dramatic shadow",
            "size": "large, impactful",
            "effect": "cinematic text styling",
            "description": "Dramatic movie subtitle style"
        },
        "soviet_cartoon": {
            "position": "bottom",
            "background": "hand-drawn banner effect",
            "font": "playful, rounded",
            "color": "warm colors (cream, yellow)",
            "size": "medium",
            "effect": "hand-drawn look",
            "description": "Classic Soviet cartoon title card style"
        }
    }

    return styles.get(visual_style, styles["documentary"])


def build_frame_prompt(
    scene: Dict,
    frame_type: str,
    script: Dict,
    refs: List[Tuple[str, Path]],
    caption_style: Dict,
    is_chained: bool = False  # True if we're using frame_first as reference for frame_last
) -> str:
    """Build NanoBanana prompt for a scene frame with all references."""
    scene_type = scene.get("type", "")
    method = scene.get("method", "veo3")
    visual_style = script.get("visual_style", "documentary")
    is_v2 = script.get("format_version") == "v2"

    # V2 format: single frame per scene with "frame" and "caption" fields
    if is_v2:
        caption = scene.get("caption", "")
        frame_desc = scene.get("frame", "")
        emotion = scene.get("emotion", "")
        if emotion:
            frame_desc = f"{frame_desc}. Emotional tone: {emotion}"
    else:
        # V1 format: first/last frames with caption_first/caption_last
        frames = scene.get("frames", {})

        # Get caption
        if frame_type == "first":
            caption = scene.get("caption_first", scene.get("caption", ""))
        elif frame_type == "last":
            caption = scene.get("caption_last", scene.get("caption", ""))
        else:
            caption = scene.get("caption", "")

        # Get frame description
        # For kenburns scenes that only have "image", use it for both first/last
        frame_desc = frames.get(frame_type, "")
        if not frame_desc and method == "kenburns":
            frame_desc = frames.get("image", "")

    # Get action (what happens during the scene)
    action = scene.get("action", "")

    # Get object placement for this specific frame
    object_placement = scene.get("object_placement", {}).get(frame_type, {})

    # Build object placement text
    placement_text = ""
    if object_placement:
        placement_lines = []
        for obj_id, placement in object_placement.items():
            obj_info = script.get("objects", {}).get(obj_id, {})
            obj_desc = obj_info.get("description", obj_id)
            placement_lines.append(f"- {obj_id} ({obj_desc}): {placement}")
        if placement_lines:
            placement_text = "\n\nOBJECT POSITIONS IN THIS FRAME (CRITICAL!):\n" + "\n".join(placement_lines)

    # Get style hints
    if visual_style:
        style_hints_dict = get_frame_style_hints(visual_style)
        style_config = get_style_config(visual_style)
        style_hints = f"""VISUAL STYLE: {style_config.get('name', visual_style)}
{style_config.get('prompt_hint', '')}
Lighting: {style_hints_dict.get('lighting', 'professional lighting')}
Color grade: {style_hints_dict.get('color_grade', 'natural colors')}
Atmosphere: {style_hints_dict.get('atmosphere', 'neutral')}"""
    else:
        style_hints = "Professional photography, natural lighting."

    # Build reference descriptions
    ref_descriptions = []
    for label, path in refs:
        ref_type, ref_id = label.split(":", 1)
        if ref_type == "character":
            if ref_id == "storyteller":
                info = script.get("storyteller", {})
            else:
                info = script.get("characters", {}).get(ref_id, {})
            ref_descriptions.append(f"CHARACTER '{ref_id}': {info.get('description', '')}")
        elif ref_type == "object":
            info = script.get("objects", {}).get(ref_id, {})
            ref_descriptions.append(f"OBJECT '{ref_id}': {info.get('description', '')}")
        elif ref_type == "location":
            info = script.get("locations", {}).get(ref_id, {})
            ref_descriptions.append(f"LOCATION '{ref_id}': {info.get('description', '')}")

    refs_text = "\n".join(ref_descriptions) if ref_descriptions else "No specific references"

    # Build caption instruction with consistent style
    caption_instruction = ""
    if caption:
        caption_instruction = f"""
CAPTION TO RENDER (CRITICAL - USE CONSISTENT STYLE!):
Text: "{caption}"
Style: {caption_style['description']}
- Position: {caption_style['position']}
- Background: {caption_style['background']}
- Font: {caption_style['font']}
- Color: {caption_style['color']}
- Size: {caption_style['size']}
- Effect: {caption_style['effect']}

IMPORTANT: This caption style MUST be identical across ALL frames in this video!"""

    # Chain instruction for frame_last
    chain_instruction = ""
    if is_chained:
        chain_instruction = """
CONTINUITY - CHANGE ONLY WHAT THE ACTION DESCRIBES:
The FIRST reference image is the FIRST FRAME of this same scene.

LOGICAL CONSISTENCY:
- If action is "son hands book to father" → ONLY the book moves between hands
- If action is "father walks to shelf" → ONLY father's position changes
- If action is "son sits down" → ONLY son goes from standing to sitting
- Everything NOT mentioned in action stays EXACTLY the same!

ALWAYS KEEP:
- Same character faces (same person!)
- Same clothes
- Same room layout and furniture positions
- Same lighting and camera angle

The frame_last must be a LOGICAL RESULT of the described action.
Don't randomly change things that weren't part of the action!"""

    # Build action context
    action_text = ""
    if action and frame_type == "last":
        action_text = f"""

SCENE ACTION (what happened between frames):
{action}

CRITICAL - THIS IS THE END STATE!
The frame_last must look like a COMPLETELY DIFFERENT PHOTOGRAPH from frame_first:
- Different body poses (sitting→standing, arms down→arms up, etc.)
- Different positions in the room (moved closer/further, turned around, etc.)
- Different facial expressions (neutral→surprised, calm→emotional, etc.)
- Objects in different hands/positions as the action describes

If the first frame shows someone reaching for something, this frame shows them HOLDING it.
If the first frame shows neutral expression, this frame shows REACTION to what happened.
Make it OBVIOUS that time has passed and action occurred!"""
    elif action and frame_type == "first":
        action_text = f"\n\nSCENE ACTION (will happen next):\n{action}\nThis frame shows the BEGINNING state - before the action starts. Pose and expression should be INITIAL state."

    # Build full prompt
    prompt = f"""Generate a video frame based on the reference images provided.

SCENE TYPE: {scene_type}
FRAME: {frame_type} ({"start" if frame_type == "first" else "end"} of scene)

{style_hints}

FRAME DESCRIPTION:
{frame_desc}{action_text}{placement_text}

REFERENCE IMAGES PROVIDED (use ALL of them for consistency):
{refs_text}
{chain_instruction}
{caption_instruction}

CRITICAL REQUIREMENTS:
- ALL text MUST be in RUSSIAN (Cyrillic script)
- Match reference images for character IDENTITY (same face, same clothes)
- Caption style MUST be consistent (see style above)
- Objects MUST be positioned EXACTLY as specified in OBJECT POSITIONS section!
- First and last frames must be like TWO DIFFERENT PHOTOS - dramatically different poses!

VERTICAL COMPOSITION (9:16) - ABSOLUTELY CRITICAL:
- Fill the ENTIRE vertical canvas edge-to-edge
- ZERO black bars on top or bottom
- ZERO blur bars or letterboxing
- ZERO padding of any kind
- Content must reach ALL FOUR EDGES of the frame
- Frame like Instagram Stories / TikTok - full bleed vertical

If you add ANY black/blur bars, the image is REJECTED. Fill the entire 9:16 canvas!

Output a single cohesive image combining all elements."""

    return prompt


def generate_scene_frames(
    scene: Dict,
    scene_dir: Path,
    sheets_dir: Path,
    api_key: str,
    script: Dict,
    caption_style: Dict,
    aspect_ratio: str = "9:16",
    skip_existing: bool = False
) -> bool:
    """Generate frames for a single scene with all references."""
    scene_id = scene["scene_id"]
    scene_type = scene.get("type", "unknown")
    method = scene.get("method", "veo3")

    print(f"\n=== Scene {scene_id}: {scene_type} ({method}) ===")

    # Check for v2 format (single frame per scene)
    is_v2 = script.get("format_version") == "v2"

    # Determine which frames to generate
    if is_v2:
        # V2 format: single frame per scene
        frames_to_generate = ["frame"]
    else:
        # V1 format: first and last frames
        frames_to_generate = ["first", "last"]

    # Gather ALL references for this scene
    refs = gather_scene_references(scene, script, sheets_dir)
    print(f"  References: {[label for label, _ in refs]}")

    if not refs:
        print(f"  WARNING: No reference sheets found for this scene!")

    # Load all reference images
    ref_images = []
    for label, path in refs:
        try:
            img_data, mime = load_image(path)
            ref_images.append((img_data, mime))
            print(f"    Loaded: {path.name}")
        except Exception as e:
            print(f"    Failed to load {path.name}: {e}")

    success = True
    frame_first_data = None  # For chaining to frame_last

    for frame_type in frames_to_generate:
        # V2 format saves as "frame.jpg", V1 as "frame_first.jpg"/"frame_last.jpg"
        if is_v2:
            output_path = scene_dir / "frame.jpg"
        else:
            output_path = scene_dir / f"frame_{frame_type}.jpg"

        if skip_existing and output_path.exists():
            print(f"  {output_path.name} already exists, skipping")
            if frame_type == "first":
                frame_first_data, _ = load_image(output_path)
            continue

        print(f"  Generating {output_path.name}...")

        # For frame_last (v1), add frame_first as additional reference
        # For v2 format, no chaining needed (single frame per scene)
        images_to_use = list(ref_images)
        is_chained = False

        if not is_v2 and frame_type == "last" and frame_first_data:
            # V1: Insert frame_first at the beginning (most important reference)
            images_to_use.insert(0, (frame_first_data, "image/jpeg"))
            is_chained = True
            print(f"    + Using frame_first as chain reference")

        # Build prompt
        prompt = build_frame_prompt(scene, frame_type, script, refs, caption_style, is_chained)

        # Save prompt for debugging
        prompt_name = "nb_prompt.txt" if is_v2 else f"nb_prompt_{frame_type}.txt"
        prompt_path = scene_dir / prompt_name
        with open(prompt_path, "w") as f:
            f.write(prompt)
            f.write(f"\n\n# References used: {[label for label, _ in refs]}")
            f.write(f"\n# Format: {'v2' if is_v2 else 'v1'}")
            if is_chained:
                f.write("\n# Chained from frame_first")

        # Generate image
        result_bytes, result_mime = call_nanobanana(
            api_key=api_key,
            prompt=prompt,
            images=images_to_use if images_to_use else None,
            aspect_ratio=aspect_ratio,
        )

        if result_bytes:
            with open(output_path, "wb") as f:
                f.write(result_bytes)
            size_kb = len(result_bytes) / 1024
            print(f"    Saved: {output_path.name} ({size_kb:.1f} KB)")

            # Store for chaining (v1 only)
            if not is_v2 and frame_type == "first":
                frame_first_data = result_bytes
        else:
            print(f"    FAILED to generate {output_path.name}")
            success = False

        time.sleep(2)  # Rate limiting

    return success


def main():
    parser = argparse.ArgumentParser(description="Generate all frames for microdrama")
    parser.add_argument("task_dir", help="Task directory path")
    parser.add_argument("--script", default="script.json", help="Script filename")
    parser.add_argument("--scene", "-s", type=int, help="Generate only this scene number")
    parser.add_argument("--aspect", "-a", default="9:16", help="Aspect ratio")
    parser.add_argument("--skip-existing", "-k", action="store_true", help="Skip existing frames")

    args = parser.parse_args()

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

    # Find script file
    script_path = task_dir / "script" / args.script
    if not script_path.exists():
        for name in ["script.json", "script_v6.json", "script_v5.json"]:
            alt_path = task_dir / "script" / name
            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: {script_path.name}")
    with open(script_path) as f:
        script = json.load(f)

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

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

    aspect_ratio = script.get("aspect_ratio", args.aspect)
    visual_style = script.get("visual_style", "documentary")

    # Get consistent caption style for entire video
    caption_style = get_caption_style(script)

    print(f"\n{'='*60}")
    print(f"Frame Generation (Enhanced)")
    print(f"Task: {task_dir.name}")
    print(f"Scenes: {len(scenes)}")
    print(f"Aspect ratio: {aspect_ratio}")
    print(f"Visual style: {visual_style}")
    print(f"Caption style: {caption_style['description']}")
    print(f"{'='*60}")

    # Check for sheets directory
    if not sheets_dir.exists():
        print(f"\nWARNING: Sheets directory not found: {sheets_dir}")
        print("Run download_references.py first to generate sheets")
    else:
        sheets = list(sheets_dir.glob("*.jpg")) + list(sheets_dir.glob("*.png"))
        print(f"Available sheets: {len(sheets)}")

    success_count = 0
    fail_count = 0

    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)

        if generate_scene_frames(
            scene, scene_dir, sheets_dir, api_key, script,
            caption_style, aspect_ratio, args.skip_existing
        ):
            success_count += 1
        else:
            fail_count += 1

    print(f"\n{'='*60}")
    print(f"RESULTS: {success_count} success, {fail_count} failed")
    print(f"{'='*60}\n")

    if fail_count > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()
