#!/usr/bin/env python3
"""
Generate scene videos using Veo3 (async) via LaoZhang API.
Falls back to Kling O1 via Pollo.ai if Veo3 fails.

Uses first/last frames from each scene for interpolation.
Runs all scenes automatically without approval prompts.
"""

import os
import sys
import json
import requests
import time
from pathlib import Path
from typing import Optional

# Add parent scripts directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))

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

# Use async Veo3 instead of sync
from laozhang_veo3_async import call_veo3_laozhang_async, poll_veo3_laozhang_async, download_video_laozhang
# Kling O1 fallback via Pollo.ai
from pollo_proxy import call_kling_o1_pollo, poll_kling_pollo, download_video_pollo


def download_video(url: str, output_path: Path, timeout: int = 120) -> bool:
    """Download video from URL."""
    try:
        print(f"  Downloading video...")
        response = requests.get(url, timeout=timeout, stream=True)
        response.raise_for_status()

        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        size_mb = output_path.stat().st_size / (1024 * 1024)
        print(f"  Saved: {output_path.name} ({size_mb:.1f} MB)")
        return True

    except Exception as e:
        print(f"  Download failed: {e}")
        return False


def build_video_prompt(scene: dict, style: str = "parody_meme", visual_style: str = "", transition: str = "") -> str:
    """
    Build Veo3 prompt from scene data.

    Key improvement: Describe the ACTION happening in the scene,
    not just the start/end states. Video generators need explicit
    motion/action descriptions for 8-second clips.

    For v2 format, transition specifies how to move between the two frames.
    """
    scene_type = scene.get("type", "")
    duration = scene.get("duration_sec", 8)

    # Transition hints for v2 format
    transition_hints = {
        "cut": "Quick cut transition - abrupt scene change",
        "dissolve": "Smooth dissolve transition - gradual blend between scenes",
        "zoom_in": "Camera zooms in during transition",
        "zoom_out": "Camera zooms out during transition, revealing more",
        "morph": "Morphing transition - elements transform smoothly",
        "fade_out": "Fade to black then fade in to next scene",
        "wipe": "Wipe transition across the frame"
    }
    transition_context = transition_hints.get(transition, "") if transition else ""

    # Get frame descriptions
    frames = scene.get("frames", {})
    first_desc = frames.get("first", "")
    last_desc = frames.get("last", "")

    # Get captions that describe the narrative
    caption_first = scene.get("caption_first", "")
    caption_last = scene.get("caption_last", "")

    # Get sound hints from scene (if available)
    scene_sound_hints = scene.get("sound_hints", "")

    # Build style context based on visual_style
    style_hints = ""
    if visual_style == "home_video":
        style_hints = "VHS home video aesthetic, nostalgic 90s film grain, warm faded colors, authentic amateur footage feel"
    elif visual_style == "documentary":
        style_hints = "Documentary style, BBC/Discovery aesthetic, professional but intimate"
    elif "sepia" in visual_style.lower():
        style_hints = "SEPIA VINTAGE TONE, period-appropriate film grain texture"
    else:
        style_hints = "Cinematic photorealistic style, natural lighting"

    # Use scene's ACTION field directly if available (preferred!)
    # Only fall back to generated description if no action specified
    scene_action = scene.get("action", "")
    if scene_action:
        # Use the script's action directly - it should describe physical movement
        action_description = scene_action
    else:
        # Fallback: try to infer action from frame descriptions
        action_description = _generate_action_description(
            scene_type=scene_type,
            first_frame=first_desc,
            last_frame=last_desc,
            caption_first=caption_first,
            caption_last=caption_last
        )

    # Camera movement based on scene type
    camera_hints = {
        "hook": "Slow dramatic push-in, building tension",
        "setup": "Steady reveal, smooth tilt or pan",
        "development": "Natural flowing movement, following the action",
        "climax": "Dynamic movement, heightened tension",
        "punchline": "Hold on reaction, subtle movement",
        "reaction": "Quick zoom to face, hold on expression",
        "reveal": "Dynamic zoom out, dramatic reveal",
        "flashback": "Slow dreamy movement, nostalgic",
        "twist": "Dramatic slow motion, theatrical"
    }.get(scene_type, "Smooth subtle movement")

    # Sound hints - use scene-specific if available, otherwise defaults
    if scene_sound_hints:
        sound_context = scene_sound_hints
    else:
        sound_context = {
            "hook": "ambient sounds, anticipation",
            "setup": "natural environment sounds, footsteps",
            "development": "ambient sounds, subtle reactions",
            "climax": "dramatic atmosphere, emotional sounds",
            "punchline": "beat of silence, then knowing exhale",
            "reaction": "sharp gasp, shocked inhale",
            "reveal": "dramatic whoosh, impact sound"
        }.get(scene_type, "ambient sounds, reactions")

    # Build transition section if specified
    transition_section = ""
    if transition_context:
        transition_section = f"\nTRANSITION: {transition_context}\n"

    prompt = f"""Generate an 8-second video clip.

VISUAL STYLE: {style_hints}

ACTION (what happens in this 8-second clip):
{action_description}

STARTING FRAME: {first_desc}
ENDING FRAME: {last_desc}
{transition_section}
CAMERA: {camera_hints}

AUDIO: Include natural sound effects:
- {sound_context}
- Character reactions: sighs, exhales, gasps (NO speech, NO dialogue)
- NO background music
- Only diegetic sounds (sounds that exist in the scene)

IMPORTANT:
- Photorealistic, natural movements
- Continuous smooth motion throughout 8 seconds
- Keep any text/captions visible that appear in the frames
- Subtle expressions and gestures, not over-dramatic

Output 8-second video with natural motion and sound effects."""

    return prompt


def _generate_action_description(
    scene_type: str,
    first_frame: str,
    last_frame: str,
    caption_first: str = "",
    caption_last: str = ""
) -> str:
    """
    Generate a description of what ACTION happens in the scene.

    This bridges the gap between static frame descriptions and
    actual motion that should occur over 8 seconds.
    """
    # Build context from captions (they tell the story)
    story_context = ""
    if caption_first or caption_last:
        story_context = f"Story context: '{caption_first}' → '{caption_last}'"

    # Analyze the transition to generate action
    # Look for key verbs and changes between frames
    action_lines = []

    # Detect character movement patterns
    if "holding" in first_frame.lower() and "pulling out" in last_frame.lower():
        action_lines.append("Character reaches into the object, slowly pulls something out, reveals it")
    elif "approaching" in first_frame.lower() or "stands" in first_frame.lower():
        action_lines.append("Character moves toward the subject, natural walking motion")
    elif "examining" in last_frame.lower() or "looking" in last_frame.lower():
        action_lines.append("Character's gaze shifts, eyes examining, subtle head movement")
    elif "hands" in first_frame.lower() and "hands" in last_frame.lower():
        action_lines.append("Hands move deliberately, fingers interact with objects")

    # Detect emotional transitions
    if "thoughtful" in first_frame.lower() or "melancholic" in last_frame.lower():
        action_lines.append("Expression gradually shifts to show deeper emotion, slight furrowing of brow")
    elif "smile" in last_frame.lower() and "smile" not in first_frame.lower():
        action_lines.append("Expression slowly changes to a smile, eyes crinkle slightly")
    elif "realization" in last_frame.lower():
        action_lines.append("Moment of understanding dawns on face, subtle shift in expression")

    # Detect object interactions
    if "book" in first_frame.lower() or "book" in last_frame.lower():
        action_lines.append("Book pages may rustle slightly, fingers touch the cover")
    if "money" in last_frame.lower() or "banknote" in last_frame.lower() or "ruble" in last_frame.lower():
        action_lines.append("Paper currency is revealed, light catches the surface")
    if "handing" in last_frame.lower() or "taking" in last_frame.lower():
        action_lines.append("Object passes between hands, deliberate transfer motion")

    # Scene type specific actions
    type_actions = {
        "hook": "Attention-grabbing moment, character engages with camera/viewer",
        "setup": "Scene establishes, characters position themselves, environment revealed",
        "development": "Action progresses naturally, story unfolds",
        "climax": "Key moment of the scene, most significant action occurs",
        "punchline": "Reaction lands, timing pause, then subtle response"
    }

    if scene_type in type_actions and not action_lines:
        action_lines.append(type_actions[scene_type])

    # Combine all action elements
    if story_context:
        action_lines.insert(0, story_context)

    if not action_lines:
        # Fallback: describe transition between frames
        action_lines.append(f"Smooth transition from initial state to final state over 8 seconds")
        action_lines.append("Natural breathing, subtle body movement, environmental details")

    return "\n".join(f"- {line}" for line in action_lines)


def generate_scene_video(
    scene: dict,
    scene_dir: Path,
    api_key: str,
    pollo_key: str,
    model: str = "veo-3.1-fl",
    visual_style: str = "",
    objects: dict = None,
    characters: dict = None,
    first_frame_override: Path = None,
    last_frame_override: Path = None
) -> bool:
    """
    Generate video for a single scene.

    Uses async Veo3 as primary, falls back to Kling O1 if Veo3 fails.

    For v2 format, first_frame_override and last_frame_override can be passed
    to use frames from different scene directories (for scene pairing).
    """
    scene_id = scene["scene_id"]

    # Find first/last frames
    first_frame = first_frame_override
    last_frame = last_frame_override

    # If not overridden, look in scene_dir (v1 format)
    if not first_frame or not last_frame:
        for ext in ["jpg", "png"]:
            # V2 format: single frame.jpg
            single_frame = scene_dir / f"frame.{ext}"
            if single_frame.exists() and not first_frame:
                first_frame = single_frame
            if single_frame.exists() and not last_frame:
                last_frame = single_frame

            # V1 format: frame_first.jpg and frame_last.jpg
            ff = scene_dir / f"frame_first.{ext}"
            lf = scene_dir / f"frame_last.{ext}"
            if ff.exists():
                first_frame = ff
            if lf.exists():
                last_frame = lf

    if not first_frame or not last_frame:
        print(f"  ERROR: Missing frames in {scene_dir}")
        return False

    # Build prompt with object/character context
    transition = scene.get("transition", "")
    prompt = build_video_prompt(scene, visual_style=visual_style, transition=transition)

    # Add object context if available (for visual consistency hints)
    if objects:
        scene_characters = scene.get("characters", [])
        object_hints = []
        for obj_key, obj_data in objects.items():
            obj_desc = obj_data.get("description", "")
            if obj_desc:
                object_hints.append(f"- {obj_key}: {obj_desc}")
        if object_hints:
            prompt += f"\n\nOBJECT CONSISTENCY (maintain these visual details):\n" + "\n".join(object_hints)

    # Save prompt for debugging
    prompt_path = scene_dir / "veo3_prompt.txt"
    with open(prompt_path, "w") as f:
        f.write(prompt)

    print(f"  First frame: {first_frame.name}")
    print(f"  Last frame: {last_frame.name}")

    # Try Veo3 async first
    veo3_success, veo3_error = _try_veo3_async(
        prompt=prompt,
        first_frame=first_frame,
        last_frame=last_frame,
        scene_dir=scene_dir,
        api_key=api_key,
        model=model
    )

    if veo3_success:
        return True

    # Veo3 failed - try Kling O1 fallback
    print(f"  Veo3 failed: {veo3_error}")
    print(f"  Trying Kling O1 fallback...")

    if not pollo_key:
        print(f"  ERROR: No POLLO_KEY for Kling fallback")
        return False

    kling_success, kling_error = _try_kling_o1(
        prompt=prompt,
        first_frame=first_frame,
        last_frame=last_frame,
        scene_dir=scene_dir,
        api_key=pollo_key
    )

    if kling_success:
        print(f"  Kling O1 fallback SUCCEEDED")
        return True

    print(f"  Kling O1 also failed: {kling_error}")
    # Save combined error
    error_path = scene_dir / "video_error.txt"
    with open(error_path, "w") as f:
        f.write(f"Veo3 error: {veo3_error}\nKling O1 error: {kling_error}")
    return False


def _try_veo3_async(
    prompt: str,
    first_frame: Path,
    last_frame: Path,
    scene_dir: Path,
    api_key: str,
    model: str = "veo-3.1-fl"
) -> tuple:
    """
    Try generating video with async Veo3.

    Returns: (success: bool, error_message: str or None)
    """
    try:
        print(f"  Trying Veo3 async ({model})...")

        # Create task
        task_result = call_veo3_laozhang_async(
            prompt=prompt,
            api_key=api_key,
            first_frame_path=str(first_frame),
            last_frame_path=str(last_frame),
            model=model
        )

        task_id = task_result.get("task_id")
        if not task_id:
            return False, "No task_id in response"

        # Poll for completion
        result = poll_veo3_laozhang_async(
            task_id=task_id,
            api_key=api_key,
            poll_interval=10,
            timeout=600
        )

        # Save response
        response_path = scene_dir / "veo3_response.json"
        with open(response_path, "w") as f:
            json.dump(result, f, indent=2)

        # Download video
        video_url = result.get("uri")
        if video_url:
            video_path = scene_dir / "clip.mp4"
            download_video_laozhang(video_url, str(video_path))
            size_mb = video_path.stat().st_size / (1024 * 1024)
            print(f"  Veo3 video saved: {video_path.name} ({size_mb:.1f} MB)")
            return True, None

        return False, "No video URL in response"

    except Exception as e:
        error_msg = str(e)
        # Save error for debugging
        error_path = scene_dir / "veo3_error.txt"
        with open(error_path, "w") as f:
            f.write(error_msg)
        return False, error_msg


def _try_kling_o1(
    prompt: str,
    first_frame: Path,
    last_frame: Path,
    scene_dir: Path,
    api_key: str
) -> tuple:
    """
    Try generating video with Kling O1 (fallback).

    Returns: (success: bool, error_message: str or None)
    """
    try:
        print(f"  Trying Kling O1 (5 seconds, 9:16)...")

        # Create task - use 5 seconds (Kling only supports 5 or 10 with first/last frame)
        task_result = call_kling_o1_pollo(
            prompt=prompt,
            api_key=api_key,
            first_frame_path=str(first_frame),
            last_frame_path=str(last_frame),
            aspect_ratio="9:16",
            length=5  # 5 seconds - Kling with first/last frame only supports 5 or 10
        )

        task_id = task_result.get("task_id")
        if not task_id:
            return False, "No task_id in response"

        # Poll for completion
        result = poll_kling_pollo(
            task_id=task_id,
            api_key=api_key,
            poll_interval=10,
            timeout=600
        )

        # Save response
        response_path = scene_dir / "kling_response.json"
        with open(response_path, "w") as f:
            json.dump(result, f, indent=2)

        # Download video
        video_url = result.get("uri")
        if video_url:
            # Download to kling_clip.mp4 first (different resolution)
            kling_path = scene_dir / "kling_clip.mp4"
            download_video_pollo(video_url, str(kling_path))
            size_mb = kling_path.stat().st_size / (1024 * 1024)
            print(f"  Kling video downloaded: {kling_path.name} ({size_mb:.1f} MB)")

            # Normalize to 720x1280 for consistency with Veo3
            clip_path = scene_dir / "clip.mp4"
            _normalize_kling_video(kling_path, clip_path)

            return True, None

        return False, "No video URL in response"

    except Exception as e:
        error_msg = str(e)
        # Save error for debugging
        error_path = scene_dir / "kling_error.txt"
        with open(error_path, "w") as f:
            f.write(error_msg)
        return False, error_msg


def _normalize_kling_video(input_path: Path, output_path: Path):
    """
    Normalize Kling video to match Veo3 format:
    - Resolution: 720x1280 (9:16)
    - Add silent audio track (Kling has no audio)
    """
    import subprocess

    print(f"  Normalizing Kling video to 720x1280 with silent audio...")

    # FFmpeg command:
    # 1. Scale to 720x1280, preserving aspect ratio with padding
    # 2. Add silent stereo audio track
    cmd = [
        "ffmpeg", "-y",
        "-i", str(input_path),
        "-f", "lavfi", "-i", "anullsrc=channel_layout=stereo:sample_rate=44100",
        "-vf", "scale=720:1280:force_original_aspect_ratio=decrease,pad=720:1280:(ow-iw)/2:(oh-ih)/2",
        "-c:v", "libx264", "-preset", "fast", "-crf", "18",
        "-c:a", "aac", "-shortest",
        str(output_path)
    ]

    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"  WARNING: FFmpeg normalization failed: {result.stderr[:200]}")
        # Fallback: just copy the file
        import shutil
        shutil.copy(input_path, output_path)
    else:
        size_mb = output_path.stat().st_size / (1024 * 1024)
        print(f"  Normalized: {output_path.name} ({size_mb:.1f} MB)")


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Generate scene videos with Veo3 (async) + Kling fallback")
    parser.add_argument("task_dir", help="Task directory path")
    parser.add_argument("--scene", "-s", type=int, help="Generate only this scene number")
    parser.add_argument("--model", "-m", default="veo-3.1-fl",
                       help="Veo3 model (default: veo-3.1-fl)")
    parser.add_argument("--skip-existing", "-k", action="store_true",
                       help="Skip scenes that already have clip.mp4")
    parser.add_argument("--kling-only", action="store_true",
                       help="Skip Veo3, use only Kling O1")
    args = parser.parse_args()

    task_dir = Path(args.task_dir)

    # Load script (prefer v5, then v4, then v3, then default)
    script_path = task_dir / "script" / "script_v5.json"
    if not script_path.exists():
        script_path = task_dir / "script" / "script_v4.json"
    if not script_path.exists():
        script_path = task_dir / "script" / "script_v3.json"
    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)

    # API keys
    api_key = os.getenv("LAOZHANG_KEY")
    pollo_key = os.getenv("POLLO_KEY")

    if not api_key and not args.kling_only:
        print("Error: Set LAOZHANG_KEY environment variable")
        sys.exit(1)

    if not pollo_key:
        print("Warning: POLLO_KEY not set, Kling fallback will not be available")

    # Extract script metadata
    visual_style = script.get("visual_style", "")
    objects = script.get("objects", {})
    characters = script.get("characters", {})

    scenes = script.get("scenes", [])

    # Check for v2 format (single frame per scene, pairs combined into clips)
    is_v2 = script.get("format_version") == "v2"

    if args.scene:
        scenes = [s for s in scenes if s["scene_id"] == args.scene]

    print(f"\n{'='*60}")
    print(f"Video Generation (Veo3 async + Kling O1 fallback)")
    print(f"Task: {task_dir.name}")
    print(f"Model: {args.model}")
    print(f"Visual style: {visual_style}")
    print(f"Format: {'v2 (scene pairs)' if is_v2 else 'v1 (first/last frames)'}")
    print(f"Scenes: {len(scenes)}")
    if is_v2:
        print(f"Video clips to generate: {len(scenes) // 2}")
    print(f"Kling fallback: {'available' if pollo_key else 'NOT AVAILABLE'}")
    print(f"{'='*60}\n")

    success_count = 0
    fail_count = 0
    skip_count = 0

    # V2 format: pair consecutive scenes into single clips
    if is_v2:
        for i in range(0, len(scenes) - 1, 2):
            scene_a = scenes[i]
            scene_b = scenes[i + 1]
            scene_id_a = scene_a["scene_id"]
            scene_id_b = scene_b["scene_id"]

            # Output goes to first scene's directory
            scene_dir_a = task_dir / "scenes" / f"scene_{scene_id_a:03d}"
            scene_dir_b = task_dir / "scenes" / f"scene_{scene_id_b:03d}"

            print(f"\n=== Clip: Scene {scene_id_a} → Scene {scene_id_b} ===")
            print(f"  Transition: {scene_a.get('transition_to_next', 'cut')}")

            # Find frames
            first_frame = None
            last_frame = None
            for ext in ["jpg", "png"]:
                ff = scene_dir_a / f"frame.{ext}"
                lf = scene_dir_b / f"frame.{ext}"
                if ff.exists():
                    first_frame = ff
                if lf.exists():
                    last_frame = lf

            if not first_frame or not last_frame:
                print(f"  ERROR: Missing frames")
                fail_count += 1
                continue

            # Check if already exists
            clip_path = scene_dir_a / "clip.mp4"
            if args.skip_existing and clip_path.exists():
                print(f"  Skipping (clip.mp4 exists)")
                skip_count += 1
                continue

            # Get transition type
            transition = scene_a.get("transition_to_next", "cut")

            # Build combined scene data for prompt
            combined_scene = {
                "scene_id": scene_id_a,
                "type": scene_a.get("type", ""),
                "caption_first": scene_a.get("caption", ""),
                "caption_last": scene_b.get("caption", ""),
                "frames": {
                    "first": scene_a.get("frame", ""),
                    "last": scene_b.get("frame", "")
                },
                "action": f"{scene_a.get('frame', '')} → {transition} → {scene_b.get('frame', '')}",
                "sound_hints": f"{scene_a.get('sound_hints', '')}, {scene_b.get('sound_hints', '')}",
                "characters": list(set(scene_a.get("characters", []) + scene_b.get("characters", []))),
                "objects": list(set(scene_a.get("objects", []) + scene_b.get("objects", []))),
                "duration_sec": scene_a.get("duration", 4) + scene_b.get("duration", 4),
                "transition": transition
            }

            if generate_scene_video(
                scene=combined_scene,
                scene_dir=scene_dir_a,
                api_key=api_key,
                pollo_key=pollo_key,
                model=args.model,
                visual_style=visual_style,
                objects=objects,
                characters=characters,
                first_frame_override=first_frame,
                last_frame_override=last_frame
            ):
                success_count += 1
                print(f"  SUCCESS")
            else:
                fail_count += 1
                print(f"  FAILED")

            # Rate limiting between clips
            if i + 2 < len(scenes) - 1:
                print(f"\n  Waiting 5s before next clip...")
                time.sleep(5)

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

        if fail_count > 0:
            sys.exit(1)
        return

    # V1 format: each scene has first/last frames
    for scene in scenes:
        scene_id = scene["scene_id"]
        scene_type = scene.get("type", "unknown")
        scene_method = scene.get("method", "veo3")
        scene_dir = task_dir / "scenes" / f"scene_{scene_id:03d}"

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

        # DISABLED: Ken Burns scenes now processed as veo3/kling
        # if scene_method == "kenburns":
        #     print(f"  Skipping (Ken Burns scene, use generate_kenburns.py)")
        #     skip_count += 1
        #     continue

        # Check if already exists
        clip_path = scene_dir / "clip.mp4"
        if args.skip_existing and clip_path.exists():
            print(f"  Skipping (clip.mp4 exists)")
            skip_count += 1
            continue

        if generate_scene_video(
            scene=scene,
            scene_dir=scene_dir,
            api_key=api_key,
            pollo_key=pollo_key,
            model=args.model,
            visual_style=visual_style,
            objects=objects,
            characters=characters
        ):
            success_count += 1
            print(f"  SUCCESS")
        else:
            fail_count += 1
            print(f"  FAILED")

        # Rate limiting between scenes
        if scene != scenes[-1]:
            print(f"\n  Waiting 5s before next scene...")
            time.sleep(5)

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

    if fail_count > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()
