#!/usr/bin/env python3
"""
Generate videos from frames using Veo3, Kling O1, or Kling v2.6.

Supports two modes:
1. Start+end frame method (Veo3, Kling O1) - uses first_frame and last_frame
2. Single frame method (Kling v2.6) - uses only starting frame with audio generation

Reads video prompts from JSON file generated by script adaptation stage.
"""

import os
import sys
import json
from pathlib import Path
from datetime import datetime

sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))

from dotenv import load_dotenv
from laozhang_proxy import call_veo3_laozhang, download_video_laozhang
from pollo_proxy import (
    call_kling_o1_pollo,
    call_kling_v26_pollo,
    poll_kling_pollo,
    download_video_pollo
)

load_dotenv(Path(__file__).parent.parent.parent / ".env")


def generate_clip_veo3(
    clip_config: dict,
    frames_dir: Path,
    output_dir: Path,
    api_key: str,
) -> tuple:
    """Generate video using Veo3 with start+end frame."""
    clip_id = clip_config["clip_id"]
    start_frame = frames_dir / clip_config["start_frame"]
    end_frame = frames_dir / clip_config["end_frame"]
    prompt = clip_config.get("veo3_prompt") or clip_config.get("prompt")

    print("Calling Veo3...")
    result = call_veo3_laozhang(
        prompt=prompt,
        api_key=api_key,
        first_frame_path=str(start_frame),
        last_frame_path=str(end_frame),
        model="veo-3.1-fast-fl",
        duration=8,
    )

    video_url = result.get("uri")
    if video_url:
        clip_dir = output_dir / f"clip_{clip_id:02d}"
        video_path = clip_dir / "video.mp4"
        download_video_laozhang(video_url, str(video_path))
        return video_path, video_url
    return None, None


def generate_clip_kling_o1(
    clip_config: dict,
    frames_dir: Path,
    output_dir: Path,
    api_key: str,
    duration: int = 5,
) -> tuple:
    """Generate video using Kling O1 with start+end frame."""
    clip_id = clip_config["clip_id"]
    start_frame = frames_dir / clip_config["start_frame"]
    end_frame = frames_dir / clip_config["end_frame"]
    prompt = clip_config.get("veo3_prompt") or clip_config.get("prompt")

    print("Calling Kling O1...")
    result = call_kling_o1_pollo(
        prompt=prompt,
        api_key=api_key,
        first_frame_path=str(start_frame),
        last_frame_path=str(end_frame),
        aspect_ratio="9:16",
        length=duration,
    )

    task_id = result.get("task_id")
    if not task_id:
        print(f"Error: No task_id in response")
        return None, None

    print(f"Polling for completion (task: {task_id})...")
    poll_result = poll_kling_pollo(task_id, api_key, timeout=300)

    video_url = poll_result.get("uri")
    if video_url:
        clip_dir = output_dir / f"clip_{clip_id:02d}"
        video_path = clip_dir / "video.mp4"
        download_video_pollo(video_url, str(video_path))
        return video_path, video_url
    return None, None


def generate_clip_kling_v26(
    clip_config: dict,
    frames_dir: Path,
    output_dir: Path,
    api_key: str,
    duration: int = 5,
    generate_audio: bool = True,
) -> tuple:
    """Generate video using Kling v2.6 with single frame and audio."""
    clip_id = clip_config["clip_id"]

    # Support both "frame" (single) and "start_frame" (paired) format
    frame_name = clip_config.get("frame") or clip_config.get("start_frame")
    frame_path = frames_dir / frame_name
    prompt = clip_config.get("prompt") or clip_config.get("veo3_prompt")

    print("Calling Kling v2.6...")
    result = call_kling_v26_pollo(
        prompt=prompt,
        api_key=api_key,
        first_frame_path=str(frame_path),
        length=duration,
        generate_audio=generate_audio,
    )

    task_id = result.get("task_id")
    if not task_id:
        print(f"Error: No task_id in response")
        return None, None

    print(f"Polling for completion (task: {task_id})...")
    poll_result = poll_kling_pollo(task_id, api_key, timeout=300)

    video_url = poll_result.get("uri")
    if video_url:
        clip_dir = output_dir / f"clip_{clip_id:02d}"
        video_path = clip_dir / "video.mp4"
        download_video_pollo(video_url, str(video_path))
        return video_path, video_url
    return None, None


def generate_clip(
    clip_config: dict,
    frames_dir: Path,
    output_dir: Path,
    api_key: str,
    model: str = "kling-v26",
    duration: int = 5,
    generate_audio: bool = True,
) -> Path:
    """
    Generate video for a single clip.

    Args:
        clip_config: Clip configuration from video_prompts.json
        frames_dir: Directory containing frame images
        output_dir: Base output directory
        api_key: API key (LaoZhang for Veo3, Pollo for Kling)
        model: "veo3", "kling-o1", or "kling-v26"
        duration: Duration in seconds (5 or 10 for Kling)
        generate_audio: Generate audio (Kling v2.6 only)

    Returns:
        Path to generated video or None if failed
    """
    clip_id = clip_config["clip_id"]

    # Support both formats
    frame_name = clip_config.get("frame") or clip_config.get("start_frame")
    caption = clip_config.get("caption", "")
    prompt = clip_config.get("prompt") or clip_config.get("veo3_prompt", "")

    print(f"\n=== Generating Clip {clip_id} ({model}) ===")
    if caption:
        print(f"Caption: {caption}")
    print(f"Frame: {frame_name}")

    # Verify frame exists
    frame_path = frames_dir / frame_name
    if not frame_path.exists():
        print(f"Error: Frame not found: {frame_path}")
        return None

    # Create clip output directory
    clip_dir = output_dir / f"clip_{clip_id:02d}"
    clip_dir.mkdir(parents=True, exist_ok=True)

    # Save prompt
    prompt_path = clip_dir / "prompt.txt"
    with open(prompt_path, 'w') as f:
        f.write(f"Clip: {clip_id}\n")
        f.write(f"Model: {model}\n")
        f.write(f"Caption: {caption}\n")
        f.write(f"Frame: {frame_name}\n")
        f.write(f"Duration: {duration}s\n")
        f.write(f"Generate audio: {generate_audio}\n")
        f.write(f"\n--- Prompt ---\n{prompt}\n")

    try:
        if model == "veo3":
            video_path, video_url = generate_clip_veo3(
                clip_config, frames_dir, output_dir, api_key
            )
        elif model == "kling-o1":
            video_path, video_url = generate_clip_kling_o1(
                clip_config, frames_dir, output_dir, api_key, duration
            )
        elif model == "kling-v26":
            video_path, video_url = generate_clip_kling_v26(
                clip_config, frames_dir, output_dir, api_key, duration, generate_audio
            )
        else:
            print(f"Error: Unknown model: {model}")
            return None

        if video_path:
            print(f"Video saved: {video_path}")

            # Save metadata
            meta_path = clip_dir / "metadata.json"
            with open(meta_path, 'w') as f:
                json.dump({
                    "clip_id": clip_id,
                    "model": model,
                    "caption": caption,
                    "frame": frame_name,
                    "duration": duration,
                    "generate_audio": generate_audio,
                    "video_url": video_url,
                    "timestamp": datetime.now().isoformat(),
                }, f, indent=2, ensure_ascii=False)

            return video_path
        else:
            print(f"Error: No video generated")
            return None

    except Exception as e:
        print(f"Error generating video: {e}")
        import traceback
        traceback.print_exc()
        return None


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Generate videos from frames")
    parser.add_argument("prompts_file", help="Path to video_prompts.json")
    parser.add_argument("--output", "-o", help="Output directory")
    parser.add_argument("--clip", "-c", type=int, help="Generate only this clip number")
    parser.add_argument("--model", "-m", choices=["veo3", "kling-o1", "kling-v26"], default="kling-v26",
                        help="Video model (default: kling-v26)")
    parser.add_argument("--duration", "-d", type=int, default=5, choices=[5, 10],
                        help="Video duration in seconds (default: 5)")
    parser.add_argument("--no-audio", action="store_true",
                        help="Disable audio generation (Kling v2.6 only)")
    args = parser.parse_args()

    # Get API key based on model
    if args.model == "veo3":
        api_key = os.environ.get("LAOZHANG_KEY")
        if not api_key:
            print("Error: LAOZHANG_KEY not set for Veo3")
            sys.exit(1)
    else:
        api_key = os.environ.get("POLLO_KEY")
        if not api_key:
            print("Error: POLLO_KEY not set for Kling")
            sys.exit(1)

    # Load prompts
    prompts_path = Path(args.prompts_file)
    if not prompts_path.exists():
        print(f"Error: Prompts file not found: {prompts_path}")
        sys.exit(1)

    with open(prompts_path) as f:
        config = json.load(f)

    frames_dir = Path(config.get("frames_dir", prompts_path.parent.parent / "frames" / "current"))
    if not frames_dir.exists():
        print(f"Error: Frames directory not found: {frames_dir}")
        sys.exit(1)

    # Output directory
    if args.output:
        output_dir = Path(args.output)
    else:
        output_dir = prompts_path.parent.parent.parent / "videos" / "002"

    output_dir.mkdir(parents=True, exist_ok=True)

    generate_audio = not args.no_audio

    print(f"Model: {args.model}")
    print(f"Duration: {args.duration}s")
    print(f"Generate audio: {generate_audio}")
    print(f"Frames: {frames_dir}")
    print(f"Output: {output_dir}")

    # Filter clips if specified
    clips = config["clips"]
    if args.clip:
        clips = [c for c in clips if c["clip_id"] == args.clip]
        if not clips:
            print(f"Error: Clip {args.clip} not found")
            sys.exit(1)

    # Generate videos
    results = []
    for clip_config in clips:
        video_path = generate_clip(
            clip_config, frames_dir, output_dir, api_key,
            model=args.model, duration=args.duration, generate_audio=generate_audio
        )
        results.append({
            "clip": clip_config["clip_id"],
            "caption": clip_config.get("caption", ""),
            "video": str(video_path) if video_path else None,
            "success": video_path is not None,
        })

    # Summary
    print("\n=== Summary ===")
    success = sum(1 for r in results if r['success'])
    print(f"Generated {success}/{len(results)} videos")

    for r in results:
        status = "✓" if r['success'] else "✗"
        caption = r['caption'][:40] + "..." if len(r['caption']) > 40 else r['caption']
        print(f"  {status} Clip {r['clip']}: {caption}")


if __name__ == "__main__":
    main()
