#!/usr/bin/env python3
"""
Generate characteristic shot 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 actual image.

Usage:
    generate_characteristic_image_v3.py STORY_NAME [--take TAKE] [--iteration ITER]
    generate_characteristic_image_v3.py "truck-stop-revenge"
    generate_characteristic_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_characteristic_shot(
    story_name: str,
    take: str = "current",
    iteration: str = "current",
    aspect_ratio: str = "16:9",
    image_size: str = "4K"
) -> bool:
    """
    Generate characteristic shot 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: 16:9 for reference)
        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 characteristic stage path
    stage_path = take_path / "06-characteristic"
    if not stage_path.exists():
        print(f"Error: Characteristic stage not found: {stage_path}")
        print(f"Run: python run_stage_v3.py '{story_name}' --stage characteristic")
        return False

    # Resolve iteration
    if iteration == "current":
        current_link = stage_path / "current"
        if not current_link.exists():
            print(f"Error: No current iteration in characteristic 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 characteristic")
        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()

    # Output image path
    output_image = iter_path / "characteristic_shot.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, "characteristic_image", iteration, session_id)

    print(f"\nGenerating characteristic shot:")
    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
        )

        print(f"\n✓ 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["characteristic_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"
        }

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

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

        print(f"\n{'='*60}")
        print(f"Characteristic shot generated successfully!")
        print(f"{'='*60}")
        print(f"Image: {output_image}")
        print(f"\nNext steps:")
        print(f"1. Review the image - does it match your vision?")
        print(f"2. If good → this becomes reference for grid 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 image: {error_msg}")
        log_stage_error(story_name, take_num, "characteristic_image", iteration, session_id, error_msg)
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Generate characteristic shot 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 4:3
        """
    )

    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="16:9",
        choices=["9:16", "16:9", "1:1", "4:3", "3:4"],
        help="Image aspect ratio (default: 16:9 for reference)"
    )
    parser.add_argument(
        "--size",
        default="4K",
        choices=["1K", "2K", "4K"],
        help="Image size (default: 4K)"
    )

    args = parser.parse_args()

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

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
