#!/usr/bin/env python3
"""
Generate high-quality starting frame for a scene using NanoBanana.

Takes raw_frame.jpg from grid extraction and upscales it to 1080x1920
using scene description and characteristic shot as references.

Usage:
    generate_scene_frame_v3.py STORY_NAME SCENE_NUM [--take TAKE] [--feedback "text"]
"""
import argparse
import sys
import json
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv

# Load .env from repository root
load_dotenv(Path(__file__).parent.parent.parent / ".env")

# Add lib and scripts to path
sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))

from lib import (
    get_story_path,
    get_take_path,
    ensure_dir,
    update_current_symlink,
    save_metadata
)

from nanobanana_text2img import generate_image


def generate_scene_frame(
    story_name: str,
    scene_num: int,
    take: str = "current",
    feedback: str = None
) -> bool:
    """
    Generate high-quality frame for a scene.

    Args:
        story_name: Story name
        scene_num: Scene number (1-9 for 3x3 grid)
        take: Take number or "current"
        feedback: User feedback for iteration

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

    # Get take number
    if take == "current":
        current_link = story_path / "current"
        if not current_link.is_symlink():
            print(f"Error: No current take set")
            return False
        take_num = current_link.readlink().name
    else:
        take_num = take

    # Scene directory
    scene_dir = take_path / "09-frames" / f"scene_{scene_num:02d}"
    if not scene_dir.exists():
        print(f"Error: Scene directory not found: {scene_dir}")
        return False

    # Check for raw_frame.jpg
    raw_frame_path = scene_dir / "raw_frame.jpg"
    if not raw_frame_path.exists():
        print(f"Error: Raw frame not found: {raw_frame_path}")
        return False

    # Read script-adapted to get scene description
    script_path = take_path / "08-script-adapted" / "current" / "script.json"
    if not script_path.exists():
        print(f"Error: Adapted script not found: {script_path}")
        return False

    with open(script_path) as f:
        script = json.load(f)

    # Find scene in script
    scene_data = None
    for scene in script["scenes"]:
        if scene["number"] == scene_num:
            scene_data = scene
            break

    if not scene_data:
        print(f"Error: Scene {scene_num} not found in script")
        return False

    # Get characteristic shot path
    char_shot_path = take_path / "06-characteristic" / "current" / "characteristic_shot.jpg"
    if not char_shot_path.exists():
        print(f"Warning: Characteristic shot not found: {char_shot_path}")
        char_shot_path = None

    # Get next iteration number and collect all previous feedback
    existing_iterations = []
    previous_feedback = []
    for item in scene_dir.iterdir():
        if item.is_dir() and item.name.isdigit() and len(item.name) == 3:
            iter_num = int(item.name)
            existing_iterations.append(iter_num)

            # Read feedback from this iteration if it exists
            feedback_file = item / "feedback.md"
            if feedback_file.exists():
                feedback_text = feedback_file.read_text()
                # Extract feedback content (skip markdown header)
                lines = feedback_text.split('\n')
                feedback_content = '\n'.join(lines[2:]).strip() if len(lines) > 2 else feedback_text.strip()
                if feedback_content:
                    previous_feedback.append(f"Iteration {iter_num:03d}: {feedback_content}")

    if existing_iterations:
        iteration_num = max(existing_iterations) + 1
    else:
        iteration_num = 1

    iteration = f"{iteration_num:03d}"
    iter_path = scene_dir / iteration
    ensure_dir(iter_path)

    print(f"Generating high-quality frame:")
    print(f"  Story: {story_name}")
    print(f"  Take: {take_num}")
    print(f"  Scene: {scene_num}")
    print(f"  Iteration: {iteration}")
    print(f"  Raw frame: {raw_frame_path}")
    if feedback:
        print(f"  With feedback: {feedback}")
    if char_shot_path:
        print(f"  Characteristic shot: {char_shot_path}")
    print()

    # Build NanoBanana prompt with sequence context
    style_info = script.get("style", {})
    visual_style = style_info.get("visual", "cinematic")
    tone = style_info.get("tone", "dramatic")

    # Get scene description
    what_is_shown = scene_data.get("what_is_shown", "")
    shot_type = scene_data.get("shot_type", "medium")
    frame_observations = scene_data.get("frame_observations", "")
    upscale_prompt = scene_data.get("upscale_prompt", "")

    # Get total scene count
    total_scenes = len(script["scenes"])

    # Get story context
    title = script.get("title", "Story")
    adaptation_notes = script.get("adaptation_notes", "")

    # Build character continuity info
    characters = script.get("characters", [])
    character_context = []
    for char in characters:
        char_id = char.get("id", "")
        actual_appearance = char.get("actual_appearance", "")
        appears_in = char.get("appears_in", [])
        if scene_num in appears_in and actual_appearance:
            character_context.append(f"- {char_id}: {actual_appearance}")

    # Build object continuity info
    objects = script.get("objects", [])
    object_context = []
    for obj in objects:
        obj_id = obj.get("id", "")
        state_by_scene = obj.get("state_by_scene", {})
        if str(scene_num) in state_by_scene:
            object_context.append(f"- {obj_id}: {state_by_scene[str(scene_num)]}")

    # Get previous scenes context (for continuity) with raw frames
    previous_scenes = []
    previous_frame_paths = []
    for i in range(max(1, scene_num - 2), scene_num):
        prev_scene = script["scenes"][i - 1]
        prev_action = prev_scene.get('action', '')
        prev_what_is_shown = prev_scene.get('what_is_shown', '')

        # Use raw frame from previous scene (smaller file size)
        prev_scene_dir = take_path / "09-frames" / f"scene_{i:02d}"
        prev_raw_frame = prev_scene_dir / "raw_frame.jpg"
        if prev_raw_frame.exists():
            previous_frame_paths.append(str(prev_raw_frame))
            previous_scenes.append(f"Frame {i}: {prev_action}\n  Description: {prev_what_is_shown[:150]}{'...' if len(prev_what_is_shown) > 150 else ''}")
        else:
            previous_scenes.append(f"Frame {i}: {prev_action} (raw frame not found)")

    # Calculate grid position (for 3x3 grid)
    grid_layout = "3x3"  # TODO: read from metadata
    cols = 3
    row = (scene_num - 1) // cols + 1
    col = (scene_num - 1) % cols + 1
    grid_position = f"row {row}, column {col}"

    # Build comprehensive prompt
    prompt_lines = [
        f"High-quality cinematic frame in 9:16 vertical format (1080x1920 resolution).",
        f"",
        f"COMPOSITION REFERENCE PROVIDED: Use the extracted frame as the base composition - keep same framing, character positions, and overall layout.",
        f"",
        f"SEQUENCE CONTEXT:",
        f"This is Frame {scene_num} of {total_scenes} in story: \"{title}\"",
        f"Frame position in grid: {grid_position}",
        f"",
    ]

    # Add previous frames for continuity
    if previous_scenes:
        prompt_lines.extend([
            f"PREVIOUS FRAMES (visual references provided):",
        ])
        for idx, prev in enumerate(previous_scenes, 1):
            prompt_lines.append(f"  {prev}")
        prompt_lines.append("")
        prompt_lines.append(f"Use the previous frames to maintain visual consistency (character appearance, object details, lighting style).")
        prompt_lines.append("")

    # Add character continuity
    if character_context:
        prompt_lines.extend([
            f"CHARACTER CONTINUITY (maintain same appearance):",
        ])
        for char in character_context:
            prompt_lines.append(f"  {char}")
        prompt_lines.append("")

    # Add object continuity
    if object_context:
        prompt_lines.extend([
            f"OBJECT CONTINUITY (maintain consistency):",
        ])
        for obj in object_context:
            prompt_lines.append(f"  {obj}")
        prompt_lines.append("")

    # Use upscale_prompt from adapted script if available, otherwise build from scene description
    if upscale_prompt:
        prompt_lines.extend([
            f"CURRENT FRAME - UPSCALING INSTRUCTIONS:",
            f"{upscale_prompt}",
            f"",
        ])
    else:
        # Fallback: build from scene description
        prompt_lines.extend([
            f"CURRENT FRAME - SCENE DESCRIPTION:",
            f"{what_is_shown}",
            f"",
            f"VISUAL STYLE: {visual_style}, {tone}",
            f"Shot type: {shot_type}",
            f"",
        ])

        if frame_observations:
            prompt_lines.extend([
                f"FRAME DETAILS:",
                f"{frame_observations}",
                f"",
            ])

    # Add all user feedback (from previous iterations + current)
    all_feedback = []
    if previous_feedback:
        all_feedback.extend(previous_feedback)
    if feedback:
        all_feedback.append(f"Iteration {iteration}: {feedback}")

    if all_feedback:
        prompt_lines.extend([
            f"USER FEEDBACK FROM ITERATIONS:",
        ])
        for fb in all_feedback:
            prompt_lines.append(f"  {fb}")
        prompt_lines.append("")

    prompt_lines.extend([
        f"TECHNICAL REQUIREMENTS:",
        f"- Increase resolution to 1080x1920 (9:16 aspect ratio)",
        f"- Enhance sharpness and detail significantly",
        f"- Remove any white or black border ribbons/dividers on edges",
        f"- Maintain exact composition from reference",
        f"- Enhance texture, lighting, and clarity",
        f"- Photorealistic quality, cinematic finish",
        f"",
        f"CRITICAL: Remove any divider artifacts (white/black borders) from edges. Clean, full-frame image.",
        f"",
        f"NEGATIVE PROMPT: low resolution, blurry, soft focus, white borders, black borders, divider lines, grid lines, watermark, text overlay, distorted, warped"
    ])

    prompt_text = "\n".join(prompt_lines)

    # Save prompt
    prompt_file = iter_path / "prompt.txt"
    prompt_file.write_text(prompt_text)
    print(f"✓ Saved prompt: {prompt_file}")

    if feedback:
        feedback_file = iter_path / "feedback.md"
        feedback_file.write_text(f"# User Feedback\n\n{feedback}\n")
        print(f"✓ Saved feedback: {feedback_file}")

    # Generate frame with NanoBanana
    output_path = iter_path / "frame.jpg"

    print(f"\nGenerating frame with NanoBanana...")
    print(f"  Using raw frame as composition reference")
    if previous_frame_paths:
        print(f"  Using {len(previous_frame_paths)} previous raw frame(s) for continuity")
    if char_shot_path:
        print(f"  Using characteristic shot as style reference")

    try:
        # Build additional references list (previous frames + characteristic shot)
        additional_refs = []
        additional_refs.extend(previous_frame_paths)
        if char_shot_path:
            additional_refs.append(str(char_shot_path))

        result_path = generate_image(
            prompt=prompt_text,
            output_path=str(output_path),
            aspect_ratio="9:16",
            image_size="4K",
            reference_image_path=str(raw_frame_path),
            reference_strength=0.7,  # Not used by Gemini, but kept for compatibility
            additional_reference_paths=additional_refs if additional_refs else None
        )

        if result_path and Path(result_path).exists():
            print(f"✓ Generated frame: {result_path}")

            # Create metadata
            # Make previous frame paths relative for metadata
            prev_frames_relative = []
            for prev_path in previous_frame_paths:
                # Convert absolute path to relative from this iteration directory
                prev_scene_num = int(Path(prev_path).parent.name.split('_')[1])
                prev_frames_relative.append(f"../scene_{prev_scene_num:02d}/raw_frame.jpg")

            metadata = {
                "scene": scene_num,
                "iteration": iteration,
                "timestamp": datetime.now().isoformat(),
                "source_raw_frame": "../raw_frame.jpg",
                "source_script": "../../08-script-adapted/current/script.json",
                "characteristic_shot": "../../06-characteristic/current/characteristic_shot.jpg" if char_shot_path else None,
                "model": "nanobanana-gemini3pro",
                "aspect_ratio": "9:16",
                "resolution": "1080x1920",
                "reference_images": {
                    "primary": "../raw_frame.jpg",
                    "previous_frames": prev_frames_relative,
                    "characteristic_shot": "../../06-characteristic/current/characteristic_shot.jpg" if char_shot_path else None
                },
                "prompt_file": "prompt.txt",
                "feedback": feedback,
                "previous_feedback_count": len(previous_feedback),
                "all_feedback": all_feedback if all_feedback else None,
                "used_upscale_prompt": bool(upscale_prompt),
                "upscale_prompt": upscale_prompt[:200] + "..." if upscale_prompt and len(upscale_prompt) > 200 else upscale_prompt,
                "scene_description": what_is_shown[:200] + "..." if len(what_is_shown) > 200 else what_is_shown
            }

            metadata_file = iter_path / "metadata.json"
            save_metadata(metadata_file, metadata)
            print(f"✓ Saved metadata: {metadata_file}")

            # Update current symlink
            update_current_symlink(scene_dir, iteration)
            print(f"✓ Updated current → {iteration}")

            print(f"\n{'='*60}")
            print(f"Scene {scene_num} frame generation completed!")
            print(f"{'='*60}")
            print(f"Output: {output_path}")
            print(f"\nReview the frame:")
            print(f"  - Is composition correct?")
            print(f"  - Are divider artifacts removed?")
            print(f"  - Is quality/sharpness good?")
            print(f"\nIf changes needed:")
            print(f"  python generate_scene_frame_v3.py {story_name} {scene_num} --feedback 'your feedback'")

            return True
        else:
            print(f"Error: Frame generation failed")
            return False

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


def main():
    parser = argparse.ArgumentParser(
        description="Generate high-quality starting frame for a scene"
    )
    parser.add_argument("story_name", help="Story name")
    parser.add_argument("scene_num", type=int, help="Scene number (1-9 for 3x3 grid)")
    parser.add_argument(
        "--take",
        default="current",
        help="Take number (default: current)"
    )
    parser.add_argument(
        "--feedback",
        help="User feedback for iteration"
    )

    args = parser.parse_args()

    success = generate_scene_frame(
        args.story_name,
        args.scene_num,
        args.take,
        args.feedback
    )

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
