#!/usr/bin/env python3
"""
Generate grid image using NanoBanana.

This script is part of the orchestration layer - it runs after Claude generates
the nanobanana_prompt.txt and calls the NanoBanana API to create the grid image.

Usage:
    generate_grid_image_v3.py STORY_NAME [--take TAKE] [--iteration ITER]
    generate_grid_image_v3.py "truck-stop-revenge"
    generate_grid_image_v3.py "truck-stop-revenge" --take 001 --iteration 002
"""
import argparse
import sys
import os
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv

# Load .env from repository root
repo_root = Path(__file__).parent.parent.parent
env_path = repo_root / ".env"
if env_path.exists():
    load_dotenv(env_path)

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

from nanobanana_text2img import generate_image
from lib import (
    get_story_path,
    get_take_path,
    save_metadata,
    log_stage_start,
    log_stage_complete,
    log_stage_error
)


def generate_grid(
    story_name: str,
    take: str = "current",
    iteration: str = "current",
    aspect_ratio: str = "9:16",
    image_size: str = "4K"
) -> bool:
    """
    Generate grid image from Claude's prompt.

    Args:
        story_name: Story name
        take: Take number or "current"
        iteration: Iteration number or "current"
        aspect_ratio: Image aspect ratio (default: 9:16 for vertical video)
        image_size: Image size (1K, 2K, 4K)

    Returns:
        True if successful, False otherwise
    """
    # Get paths
    try:
        story_path = get_story_path(story_name)
        take_path = get_take_path(story_name, take)
        take_num = take_path.name
    except ValueError as e:
        print(f"Error: {e}")
        return False

    # Get grid stage path
    stage_path = take_path / "07-grids"
    if not stage_path.exists():
        print(f"Error: Grid stage not found: {stage_path}")
        print(f"Run: python run_stage_v3.py '{story_name}' --stage grid")
        return False

    # Resolve iteration
    if iteration == "current":
        current_link = stage_path / "current"
        if not current_link.exists():
            print(f"Error: No current iteration in grid stage")
            return False
        iteration = current_link.readlink().name

    iter_path = stage_path / iteration
    if not iter_path.exists():
        print(f"Error: Iteration {iteration} not found")
        return False

    # Check for prompt file
    prompt_file = iter_path / "nanobanana_prompt.txt"
    if not prompt_file.exists():
        print(f"Error: Prompt file not found: {prompt_file}")
        print(f"Run: python run_stage_v3.py '{story_name}' --stage grid")
        return False

    # Read prompt
    print(f"Reading prompt from: {prompt_file}")
    prompt_text = prompt_file.read_text().strip()

    # Extract just the prompt (before metadata section)
    if "---" in prompt_text:
        prompt_text = prompt_text.split("---")[0].strip()

    # Check for generation config
    config_file = iter_path / "generation_config.json"
    reference_image_path = None
    reference_strength = 0.5

    if config_file.exists():
        import json
        config = json.loads(config_file.read_text())
        if config.get("use_reference_grid", False):
            ref_path = config.get("reference_grid_path")
            if ref_path:
                # Resolve relative path
                reference_image_path = (iter_path / ref_path).resolve()
                reference_strength = config.get("reference_strength", 0.5)
                print(f"Using reference grid: {reference_image_path}")
                print(f"Reference strength: {reference_strength}")

    # Output image path
    output_image = iter_path / "grid.png"

    # Load session ID for logging
    session_file = story_path / "session.txt"
    session_id = session_file.read_text().strip() if session_file.exists() else "unknown"

    # Log start
    log_stage_start(story_name, take_num, "grid_image", iteration, session_id)

    print(f"\nGenerating grid image:")
    print(f"  Story: {story_name}")
    print(f"  Take: {take_num}")
    print(f"  Iteration: {iteration}")
    print(f"  Aspect ratio: {aspect_ratio}")
    print(f"  Size: {image_size}")
    print(f"  Output: {output_image}")
    print()

    # Generate image
    try:
        result_path = generate_image(
            prompt=prompt_text,
            output_path=str(output_image),
            aspect_ratio=aspect_ratio,
            image_size=image_size,
            reference_image_path=str(reference_image_path) if reference_image_path else None,
            reference_strength=reference_strength
        )

        print(f"\n✓ Grid image generated: {result_path}")

        # Update metadata
        metadata_file = iter_path / "metadata.json"
        if metadata_file.exists():
            import json
            metadata = json.loads(metadata_file.read_text())
        else:
            metadata = {}

        metadata["grid_image"] = {
            "generated": True,
            "timestamp": datetime.now().isoformat(),
            "image_path": str(output_image.name),
            "aspect_ratio": aspect_ratio,
            "image_size": image_size,
            "prompt_file": str(prompt_file.name),
            "model": "nanobanana-gemini3pro",
            "reference_image": str(reference_image_path) if reference_image_path else None,
            "reference_strength": reference_strength if reference_image_path else None
        }

        save_metadata(metadata_file, metadata)
        print(f"✓ Updated metadata")

        # Log completion
        log_stage_complete(
            story_name, take_num, "grid_image",
            iteration, session_id, str(output_image)
        )

        print(f"\n{'='*60}")
        print(f"Grid image generated successfully!")
        print(f"{'='*60}")
        print(f"Image: {output_image}")
        print(f"\nNext steps:")
        print(f"1. Review the grid - do frames look consistent?")
        print(f"2. If good → extract individual frames for video generation")
        print(f"3. If needs changes → run stage again with --feedback")
        print(f"\nTo view image:")
        print(f"  open {output_image}  # macOS")
        print(f"  xdg-open {output_image}  # Linux")

        return True

    except Exception as e:
        error_msg = str(e)
        print(f"\nError generating grid: {error_msg}")
        log_stage_error(story_name, take_num, "grid_image", iteration, session_id, error_msg)
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Generate grid image using NanoBanana",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s "truck-stop-revenge"
  %(prog)s "truck-stop-revenge" --take 001 --iteration 002
  %(prog)s "truck-stop-revenge" --aspect-ratio 16:9
        """
    )

    parser.add_argument(
        "story_name",
        help="Story name"
    )
    parser.add_argument(
        "--take",
        default="current",
        help="Take number (default: current)"
    )
    parser.add_argument(
        "--iteration",
        default="current",
        help="Iteration number (default: current)"
    )
    parser.add_argument(
        "--aspect-ratio",
        default="9:16",
        choices=["9:16", "16:9", "1:1", "4:3", "3:4"],
        help="Image aspect ratio (default: 9:16 for vertical video)"
    )
    parser.add_argument(
        "--size",
        default="4K",
        choices=["1K", "2K", "4K"],
        help="Image size (default: 4K)"
    )

    args = parser.parse_args()

    success = generate_grid(
        args.story_name,
        args.take,
        args.iteration,
        args.aspect_ratio,
        args.size
    )

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
