#!/usr/bin/env python3
"""
Compare Veo 3.1 fast vs standard models.

Usage:
    python compare_veo_models.py <task_path>

Example:
    python compare_veo_models.py /mnt/video/wb_542710110/20260122_150854_lukp
"""

import json
import os
import sys
from pathlib import Path
from dotenv import load_dotenv

# Add scripts dir to path
sys.path.insert(0, str(Path(__file__).parent))

from laozhang_proxy import call_veo3_laozhang, download_video_laozhang

# Model mapping: fast -> standard
MODEL_MAPPING = {
    "veo-3.1-landscape-fast-fl": "veo-3.1-landscape-fl",
    "veo-3.1-fast-fl": "veo-3.1-fl",
}


def run_comparison(task_path: str):
    """Generate standard model video for an existing fast generation."""
    task_dir = Path(task_path)

    # Load metadata
    metadata_path = task_dir / "metadata.json"
    if not metadata_path.exists():
        print(f"Error: metadata.json not found in {task_dir}")
        return False

    with open(metadata_path) as f:
        metadata = json.load(f)

    # Get model info
    fast_model = metadata.get("mode_config", {}).get("veo3_model", "")
    standard_model = MODEL_MAPPING.get(fast_model)

    if not standard_model:
        print(f"Error: Unknown fast model '{fast_model}'")
        print(f"Known models: {list(MODEL_MAPPING.keys())}")
        return False

    print(f"Task: {metadata.get('task_id')}")
    print(f"Product: {metadata.get('product_name')}")
    print(f"Fast model: {fast_model}")
    print(f"Standard model: {standard_model}")
    print()

    # Get prompt
    veo3_prompt = metadata.get("prompts", {}).get("creative", {}).get("veo3_prompt", "")
    if not veo3_prompt:
        print("Error: veo3_prompt not found in metadata")
        return False

    # Get frames
    frame_act1 = task_dir / "processed" / "frame_act1.jpg"
    frame_act2 = task_dir / "processed" / "frame_act2.jpg"

    if not frame_act1.exists():
        # Try .png
        frame_act1 = task_dir / "processed" / "frame_act1.png"
    if not frame_act2.exists():
        frame_act2 = task_dir / "processed" / "frame_act2.png"

    print(f"Frame 1: {frame_act1} (exists: {frame_act1.exists()})")
    print(f"Frame 2: {frame_act2} (exists: {frame_act2.exists()})")
    print()

    # Check output doesn't already exist
    output_path = task_dir / "videos" / "video_standard.mp4"
    if output_path.exists():
        print(f"Standard video already exists: {output_path}")
        return True

    # Load API key
    load_dotenv(Path(__file__).parent.parent / ".env")
    api_key = os.environ.get("LAOZHANG_KEY")
    if not api_key:
        print("Error: LAOZHANG_KEY not set")
        return False

    print(f"Generating video with {standard_model}...")
    print(f"Prompt: {veo3_prompt[:200]}...")
    print()

    # Call Veo3 with standard model
    try:
        result = call_veo3_laozhang(
            prompt=veo3_prompt,
            api_key=api_key,
            first_frame_path=str(frame_act1) if frame_act1.exists() else None,
            last_frame_path=str(frame_act2) if frame_act2.exists() else None,
            model=standard_model,
            n=1
        )

        video_url = result.get("uri")
        if not video_url:
            print(f"Error: No video URL in response: {result}")
            return False

        print(f"Video URL: {video_url}")

        # Download video
        print(f"Downloading to {output_path}...")
        download_video_laozhang(video_url, str(output_path))

        print(f"Done! Standard video saved to: {output_path}")
        print(f"Fast video: {task_dir / 'videos' / 'video.mp4'}")

        return True

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


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    task_path = sys.argv[1]
    success = run_comparison(task_path)
    sys.exit(0 if success else 1)
