#!/usr/bin/env python3
"""
Test 3:4 aspect ratio for WB cards.

Steps:
1. Take existing generation
2. Generate new 3:4 frames with NanoBanana
3. Generate Veo3 video (will be 9:16 with possible letterboxing)
4. Crop video to 3:4

Usage:
    python test_3x4_format.py <task_path>
    python test_3x4_format.py /mnt/video/wb_47877202/20260126_230043_d5s1
"""

import json
import os
import subprocess
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_nanobanana_laozhang, call_veo3_laozhang, download_video_laozhang


def run_test(task_path: str):
    """Test 3:4 format 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)

    print(f"Task: {metadata.get('task_id')}")
    print(f"Product: {metadata.get('product_name')}")
    print()

    # 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

    # Get original prompts
    prompts = metadata.get("prompts", {}).get("creative", {})
    frame_act1_prompt = prompts.get("frame_act1", "")
    frame_act2_prompt = prompts.get("frame_act2", "")
    veo3_prompt = prompts.get("veo3_prompt", "")

    # Modify prompts for 3:4
    frame_act1_prompt_34 = frame_act1_prompt.replace("16:9", "3:4").replace("9:16", "3:4")
    frame_act2_prompt_34 = frame_act2_prompt.replace("16:9", "3:4").replace("9:16", "3:4")

    print("=== Original prompts ===")
    print(f"Act1: {frame_act1_prompt[:100]}...")
    print(f"Act2: {frame_act2_prompt[:100]}...")
    print()
    print("=== Modified for 3:4 ===")
    print(f"Act1: {frame_act1_prompt_34[:100]}...")
    print(f"Act2: {frame_act2_prompt_34[:100]}...")
    print()

    # Find source images
    images_dir = task_dir / "images"
    source_images = list(images_dir.glob("*.jpg")) + list(images_dir.glob("*.webp"))
    if not source_images:
        print(f"Error: No source images in {images_dir}")
        return False

    source_image = source_images[0]
    print(f"Source image: {source_image}")

    # Create output directory
    output_dir = task_dir / "test_3x4"
    output_dir.mkdir(exist_ok=True)

    # Step 1: Generate 3:4 frames
    print("\n=== Step 1: Generate 3:4 frames ===")

    with open(source_image, "rb") as f:
        image_data = f.read()

    mime_type = "image/jpeg" if source_image.suffix == ".jpg" else "image/webp"

    # Generate Act 1 frame
    print("\nGenerating Act 1 frame (3:4)...")
    try:
        act1_bytes, act1_mime = call_nanobanana_laozhang(
            image_data=image_data,
            prompt=frame_act1_prompt_34,
            api_key=api_key,
            mime_type=mime_type,
            aspect_ratio="3:4",
            image_size="2K"  # Higher resolution for better quality after cropping
        )

        act1_path = output_dir / "frame_act1_3x4.jpg"
        with open(act1_path, "wb") as f:
            f.write(act1_bytes)
        print(f"  Saved: {act1_path}")

        # Check dimensions
        result = subprocess.run(
            ["file", str(act1_path)],
            capture_output=True, text=True
        )
        print(f"  Info: {result.stdout.strip()}")

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

    # Generate Act 2 frame (using Act 1 as reference)
    print("\nGenerating Act 2 frame (3:4)...")
    try:
        act2_bytes, act2_mime = call_nanobanana_laozhang(
            image_data=image_data,
            prompt=frame_act2_prompt_34,
            api_key=api_key,
            mime_type=mime_type,
            aspect_ratio="3:4",
            image_size="2K",  # Higher resolution for better quality after cropping
            reference_image_data=act1_bytes,
            reference_mime_type=act1_mime
        )

        act2_path = output_dir / "frame_act2_3x4.jpg"
        with open(act2_path, "wb") as f:
            f.write(act2_bytes)
        print(f"  Saved: {act2_path}")

        result = subprocess.run(
            ["file", str(act2_path)],
            capture_output=True, text=True
        )
        print(f"  Info: {result.stdout.strip()}")

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

    # Step 2: Generate Veo3 video
    print("\n=== Step 2: Generate Veo3 video ===")
    print(f"Prompt: {veo3_prompt[:200]}...")

    try:
        result = call_veo3_laozhang(
            prompt=veo3_prompt,
            api_key=api_key,
            first_frame_path=str(act1_path),
            last_frame_path=str(act2_path),
            model="veo-3.1-fl",  # Standard vertical
            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
        video_raw_path = output_dir / "video_raw.mp4"
        print(f"Downloading to {video_raw_path}...")
        download_video_laozhang(video_url, str(video_raw_path))

        # Check dimensions
        result = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=width,height", "-of", "csv=s=x:p=0",
             str(video_raw_path)],
            capture_output=True, text=True
        )
        raw_dims = result.stdout.strip()
        print(f"  Raw video dimensions: {raw_dims}")

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

    # Step 3: Crop video to 3:4
    print("\n=== Step 3: Crop video to 3:4 ===")

    # Parse dimensions
    try:
        width, height = map(int, raw_dims.split("x"))
    except:
        print(f"Error parsing dimensions: {raw_dims}")
        return False

    # Calculate crop parameters
    # Target: 3:4 ratio
    # If current is 9:16 (0.5625), we need to crop height to make it 3:4 (0.75)
    target_ratio = 3 / 4
    current_ratio = width / height

    if current_ratio < target_ratio:
        # Video is too tall, crop height from TOP (not center)
        new_height = int(width / target_ratio)
        crop_filter = f"crop={width}:{new_height}:0:0"  # x=0, y=0 = top-left
    else:
        # Video is too wide, crop width (center horizontally)
        new_width = int(height * target_ratio)
        crop_filter = f"crop={new_width}:{height}"

    print(f"Current: {width}x{height} ({current_ratio:.4f})")
    print(f"Target: 3:4 ({target_ratio:.4f})")
    print(f"Crop filter: {crop_filter}")

    video_cropped_path = output_dir / "video_3x4.mp4"

    cmd = [
        "ffmpeg", "-y",
        "-i", str(video_raw_path),
        "-vf", crop_filter,
        "-c:a", "copy",
        str(video_cropped_path)
    ]

    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)

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

    # Verify final dimensions
    result = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-show_entries", "stream=width,height", "-of", "csv=s=x:p=0",
         str(video_cropped_path)],
        capture_output=True, text=True
    )
    final_dims = result.stdout.strip()
    print(f"Final video dimensions: {final_dims}")

    print("\n=== Done ===")
    print(f"Raw video: {video_raw_path}")
    print(f"Cropped video: {video_cropped_path}")
    print(f"\nTo compare:")
    print(f"  mpv {video_raw_path}")
    print(f"  mpv {video_cropped_path}")

    return True


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

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