#!/usr/bin/env python3
"""
Generate Ken Burns effect clips from static images using FFmpeg.

Supports zoom in, zoom out, pan left/right/up/down.
"""

import subprocess
import sys
from pathlib import Path
from typing import Optional


def generate_kenburns_clip(
    image_path: Path,
    output_path: Path,
    duration: float = 5.0,
    motion: str = "slow_zoom_in",
    start_scale: float = 1.0,
    end_scale: float = 1.2,
    fps: int = 30,
    width: int = 720,
    height: int = 1280
) -> bool:
    """
    Generate Ken Burns effect video from static image.

    Motion types:
    - slow_zoom_in: Zoom in slowly
    - slow_zoom_out: Zoom out slowly
    - pan_left: Pan from right to left
    - pan_right: Pan from left to right
    - pan_up: Pan from bottom to top
    - pan_down: Pan from top to bottom
    """

    total_frames = int(duration * fps)

    # Calculate zoom parameters
    # zoompan filter: zoom starts at 'zoom' and increases by 'z' each frame
    # For zoom in: start at start_scale, end at end_scale
    # For zoom out: start at end_scale, end at start_scale

    if motion == "slow_zoom_in":
        # Zoom in: scale increases
        zoom_start = start_scale
        zoom_delta = (end_scale - start_scale) / total_frames
        # Center the zoom
        x_expr = f"iw/2-(iw/zoom/2)"
        y_expr = f"ih/2-(ih/zoom/2)"
        zoom_expr = f"zoom+{zoom_delta:.6f}"
        initial_zoom = zoom_start

    elif motion == "slow_zoom_out":
        # Zoom out: scale decreases (but zoompan zoom must increase, so we invert)
        zoom_start = end_scale
        zoom_delta = (start_scale - end_scale) / total_frames
        x_expr = f"iw/2-(iw/zoom/2)"
        y_expr = f"ih/2-(ih/zoom/2)"
        zoom_expr = f"zoom+{zoom_delta:.6f}"
        initial_zoom = zoom_start

    elif motion == "pan_left":
        # Pan left: x moves from right to center
        zoom_start = start_scale if start_scale > 1.0 else 1.1
        zoom_delta = 0
        # Start at right edge, move to center
        x_expr = f"iw/2-(iw/zoom/2)-on*{(1/total_frames)*0.3:.6f}*iw"
        y_expr = f"ih/2-(ih/zoom/2)"
        zoom_expr = f"{zoom_start}"
        initial_zoom = zoom_start

    elif motion == "pan_right":
        # Pan right: x moves from left to center
        zoom_start = start_scale if start_scale > 1.0 else 1.1
        zoom_delta = 0
        x_expr = f"iw/2-(iw/zoom/2)+on*{(1/total_frames)*0.3:.6f}*iw"
        y_expr = f"ih/2-(ih/zoom/2)"
        zoom_expr = f"{zoom_start}"
        initial_zoom = zoom_start

    else:
        # Default: gentle zoom in
        zoom_start = 1.0
        zoom_delta = 0.1 / total_frames
        x_expr = f"iw/2-(iw/zoom/2)"
        y_expr = f"ih/2-(ih/zoom/2)"
        zoom_expr = f"zoom+{zoom_delta:.6f}"
        initial_zoom = zoom_start

    # Build FFmpeg command
    # zoompan filter creates video from still image with zoom/pan effect
    filter_complex = (
        f"zoompan=z='{zoom_expr}':"
        f"x='{x_expr}':"
        f"y='{y_expr}':"
        f"d={total_frames}:"
        f"s={width}x{height}:"
        f"fps={fps},"
        f"format=yuv420p"
    )

    cmd = [
        "ffmpeg", "-y",
        "-loop", "1",
        "-i", str(image_path),
        # Add silent audio track for consistency with Veo3 clips (MUST be before output options)
        "-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo",
        "-filter_complex", filter_complex,
        "-t", str(duration),
        "-map", "0:v", "-map", "1:a",  # Map video from input 0, audio from input 1
        "-c:v", "libx264",
        "-preset", "medium",
        "-crf", "18",
        "-pix_fmt", "yuv420p",
        "-shortest",
        "-c:a", "aac",
        "-b:a", "128k",
        str(output_path)
    ]

    print(f"  Generating Ken Burns clip: {motion}")
    print(f"  Duration: {duration}s, Scale: {start_scale} -> {end_scale}")

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"  Error: {result.stderr[-500:]}")
        return False

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


def main():
    import argparse
    import json

    parser = argparse.ArgumentParser(description="Generate Ken Burns clips from images")
    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("--script", default="script_v5.json", help="Script filename")
    args = parser.parse_args()

    task_dir = Path(args.task_dir)

    # Load script
    script_path = task_dir / "script" / args.script
    if not script_path.exists():
        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():
        print(f"Error: No script found")
        sys.exit(1)

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

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

    # Filter to kenburns scenes only
    kenburns_scenes = [s for s in scenes if s.get("method") == "kenburns"]

    print(f"\n{'='*60}")
    print(f"Ken Burns Video Generation")
    print(f"Task: {task_dir.name}")
    print(f"Scenes: {len(kenburns_scenes)} (kenburns method)")
    print(f"{'='*60}\n")

    success_count = 0
    fail_count = 0

    for scene in kenburns_scenes:
        scene_id = scene["scene_id"]
        scene_dir = task_dir / "scenes" / f"scene_{scene_id:03d}"

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

        # Find the source image (frame_first.jpg for kenburns scenes)
        image_path = None
        for ext in ["jpg", "png"]:
            img = scene_dir / f"frame_first.{ext}"
            if img.exists():
                image_path = img
                break

        if not image_path:
            print(f"  ERROR: No image found in {scene_dir}")
            fail_count += 1
            continue

        # Get motion parameters from script
        frames = scene.get("frames", {})
        motion = frames.get("motion", "slow_zoom_in")
        start_scale = frames.get("start_scale", 1.0)
        end_scale = frames.get("end_scale", 1.2)
        duration = scene.get("duration_sec", 5)

        output_path = scene_dir / "clip.mp4"

        if generate_kenburns_clip(
            image_path=image_path,
            output_path=output_path,
            duration=duration,
            motion=motion,
            start_scale=start_scale,
            end_scale=end_scale
        ):
            success_count += 1
            print(f"  SUCCESS")
        else:
            fail_count += 1
            print(f"  FAILED")

    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()
