#!/usr/bin/env python3
"""
Extract individual frames from NxN grid image.

Uses PIL to crop the grid into separate frames.
Supports 3x3 (9 frames) or 4x4 (16 frames) layouts.
Optionally upscales each frame using NanoBanana.
"""

import os
import sys
import base64
from pathlib import Path
from datetime import datetime

from PIL import Image
import requests

# Add scripts dir to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
from dotenv import load_dotenv

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

BASE_URL_NATIVE = "https://api.laozhang.ai/v1beta"


def extract_frames_from_grid(grid_path: Path, output_dir: Path, layout: str = "3x3") -> list:
    """
    Extract frames from NxN grid image using PIL.

    Args:
        grid_path: Path to grid image
        output_dir: Directory to save frames
        layout: Grid layout ("3x3" or "4x4")

    Returns list of frame paths.
    """
    # Parse layout
    cols, rows = map(int, layout.split('x'))

    print(f"Loading grid: {grid_path}")
    img = Image.open(grid_path)
    width, height = img.size
    print(f"Grid size: {width}x{height}")
    print(f"Layout: {layout} ({rows}x{cols} = {rows*cols} frames)")

    # Calculate frame dimensions
    frame_width = width // cols
    frame_height = height // rows
    print(f"Frame size: {frame_width}x{frame_height}")

    output_dir.mkdir(parents=True, exist_ok=True)
    frame_paths = []

    for row in range(rows):
        for col in range(cols):
            frame_num = row * cols + col + 1

            # Calculate crop box (left, upper, right, lower)
            left = col * frame_width
            upper = row * frame_height
            right = left + frame_width
            lower = upper + frame_height

            # Crop frame
            frame = img.crop((left, upper, right, lower))

            # Save frame
            frame_path = output_dir / f"frame_{frame_num:02d}.jpg"
            frame.save(frame_path, "JPEG", quality=95)
            frame_paths.append(frame_path)
            print(f"  Frame {frame_num}: {frame_path.name} ({frame.size[0]}x{frame.size[1]})")

    return frame_paths


def upscale_frame_nanobanana(frame_path: Path, api_key: str, output_path: Path) -> Path:
    """
    Upscale a single frame using NanoBanana.
    """
    print(f"Upscaling {frame_path.name}...")

    # Load frame
    with open(frame_path, 'rb') as f:
        image_data = f.read()

    image_base64 = base64.b64encode(image_data).decode('utf-8')

    # Upscale prompt
    prompt = f"Upscale this image to maximum quality 9:16 vertical format, maintaining all cinematic details, textures, and lighting. Keep exact same composition and content."

    payload = {
        "contents": [{
            "parts": [
                {"inline_data": {"mime_type": "image/jpeg", "data": image_base64}},
                {"text": prompt}
            ]
        }],
        "generationConfig": {
            "responseModalities": ["IMAGE"],
            "imageConfig": {
                "aspectRatio": "9:16",
                "imageSize": "4K"
            }
        }
    }

    response = requests.post(
        f"{BASE_URL_NATIVE}/models/gemini-3-pro-image-preview:generateContent",
        headers={
            "x-goog-api-key": api_key,
            "Content-Type": "application/json"
        },
        json=payload,
        timeout=180
    )

    if response.status_code != 200:
        print(f"  Error: {response.status_code}")
        return None

    data = response.json()
    parts = data.get("candidates", [{}])[0].get("content", {}).get("parts", [])

    for part in parts:
        if "inlineData" in part:
            image_bytes = base64.b64decode(part["inlineData"]["data"])
            with open(output_path, 'wb') as f:
                f.write(image_bytes)
            print(f"  Saved: {output_path.name} ({len(image_bytes) // 1024}KB)")
            return output_path

    print("  No image in response")
    return None


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Extract frames from NxN grid")
    parser.add_argument("grid_path", help="Path to grid image")
    parser.add_argument("--layout", "-l", default="3x3", help="Grid layout (3x3 or 4x4)")
    parser.add_argument("--upscale", action="store_true", help="Upscale frames with NanoBanana")
    parser.add_argument("--output", "-o", help="Output directory (default: same as grid)")
    args = parser.parse_args()

    grid_path = Path(args.grid_path)
    if not grid_path.exists():
        print(f"Error: Grid not found: {grid_path}")
        sys.exit(1)

    # Output directory
    if args.output:
        output_dir = Path(args.output)
    else:
        output_dir = grid_path.parent / "frames"

    # Extract frames
    frame_paths = extract_frames_from_grid(grid_path, output_dir, args.layout)
    print(f"\nExtracted {len(frame_paths)} frames to {output_dir}")

    # Optionally upscale
    if args.upscale:
        api_key = os.environ.get("LAOZHANG_KEY")
        if not api_key:
            print("Error: LAOZHANG_KEY not set for upscaling")
            sys.exit(1)

        upscaled_dir = output_dir / "upscaled"
        upscaled_dir.mkdir(exist_ok=True)

        print(f"\nUpscaling frames with NanoBanana...")
        for frame_path in frame_paths:
            upscaled_path = upscaled_dir / frame_path.name
            upscale_frame_nanobanana(frame_path, api_key, upscaled_path)

        print(f"\nUpscaled frames saved to {upscaled_dir}")


if __name__ == "__main__":
    main()
