#!/usr/bin/env python3
"""
Stage 6: Characteristic Shot

Two-step process:
1. Call Claude to create a NanoBanana prompt (nanobanana_prompt.txt)
2. Call NanoBanana API to generate the characteristic reference image

Creates 06-characteristic/XXX/nanobanana_prompt.txt + characteristic_shot.png
"""

import subprocess
import sys
import os
from pathlib import Path
from datetime import datetime, timezone

# Add parent to path for lib imports
sys.path.insert(0, str(Path(__file__).parent.parent))

# Add repo scripts dir for nanobanana import
repo_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(repo_root / "scripts"))

from lib import (
    save_metadata,
    ensure_dir,
    setup_generation_log,
    log_event,
    claude_env,
    build_reference_record
)
from lib.path_utils_v3 import get_next_iteration
from lib.feedback_history_v5 import get_feedback_history_for_claude
from lib.reference_analysis import build_analyze_prefix
import reference_loader
from reference_loader import (
    is_story_mode, get_episode_characters, select_character_refs,
    get_episode_locations, get_episode_objects, select_asset_refs,
    select_asset_refs_with_fallback,
    NANOBANANA_MAX_IMAGES,
)


def execute(take_path: Path, story_name: str, inputs: dict,
            feedback: str = None, refined_from: str = None,
            chat_mode: bool = False, reiterate: bool = False,
            scale_refs: bool = True) -> dict:
    """
    Execute Stage 6: Characteristic Shot.

    Args:
        take_path: Path to take directory
        story_name: Story name (slug)
        inputs: Parent stage iterations, e.g. {"03-script": "006", "04-assets": "004"}
        feedback: User feedback text (if any)
        refined_from: Iteration being refined (if refinement)
        chat_mode: If True, skip Claude and refine directly with NanoBanana chat
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 6: Characteristic Shot")

    # Resolve inputs
    script_iter = inputs.get("03-script")
    assets_iter = inputs.get("04-assets")

    if not script_iter:
        return {"status": "error", "error": "No 03-script iteration in inputs"}
    if not assets_iter:
        return {"status": "error", "error": "No 04-assets iteration in inputs"}

    script_file = take_path / "03-script" / script_iter / "script.json"
    assets_file = take_path / "04-assets" / assets_iter / "assets.json"

    if not script_file.exists():
        return {"status": "error", "error": f"Script file not found: {script_file}"}
    if not assets_file.exists():
        return {"status": "error", "error": f"Assets file not found: {assets_file}"}

    logger.info(f"Script: {script_file}")
    logger.info(f"Assets: {assets_file}")

    # Determine iteration number
    char_stage = take_path / "06-characteristic"
    ensure_dir(char_stage)
    iteration = get_next_iteration(char_stage)

    iter_path = char_stage / iteration
    ensure_dir(iter_path)

    logger.info(f"Creating iteration {iteration}")
    if refined_from:
        logger.info(f"Refined from: {refined_from}")

    # Save feedback if provided
    if feedback:
        (iter_path / "feedback.txt").write_text(feedback)
        logger.info("Saved feedback.txt")

    # Load session ID
    session_file = take_path.parent / "session.txt"
    if not session_file.exists():
        return {"status": "error", "error": "Session file not found"}

    session_id = session_file.read_text().strip()
    logger.info(f"Session ID: {session_id}")

    # Feedback history
    feedback_history = get_feedback_history_for_claude(
        take_path, "06-characteristic", refined_from
    )

    # Build prompt
    prompt_template_path = Path(__file__).parent.parent.parent / "prompts" / "stage_characteristic.txt"

    if not prompt_template_path.exists():
        return {"status": "error", "error": f"Prompt template not found: {prompt_template_path}"}

    prompt = prompt_template_path.read_text()

    # Add template/references context if project has them
    project_path = take_path.parent
    refs = reference_loader.load_project_references(project_path)
    template_context = reference_loader.get_full_prompt_context(project_path)
    if template_context:
        prompt += "\n\n" + template_context
        logger.info(f"Added template context ({len(template_context)} chars)")

    # Add feedback
    if feedback_history:
        prompt += f"\n\n## Previous Feedback\n\n{feedback_history}\n"
    elif feedback:
        prompt += f"\n\n## User Feedback\n\n{feedback}\n"

    # Prepend output directory
    story_path = take_path.parent
    relative_iter_path = iter_path.relative_to(story_path)
    prompt = f"## Output Directory\n\nWrite ALL output files to: {relative_iter_path}/\n\n{prompt}"

    # Determine generation mode
    if reiterate and refined_from:
        generation_mode = "reiterate"
        logger.info(f"Mode: REITERATE (same prompt from {refined_from}, new image)")
    elif chat_mode and refined_from and feedback:
        generation_mode = "chat"
        logger.info(f"Mode: CHAT (continuing conversation from {refined_from})")
    elif refined_from and feedback:
        generation_mode = "refine"
        logger.info(f"Mode: REFINE (regenerating with accumulated feedback)")
    else:
        generation_mode = "new"
        logger.info(f"Mode: NEW (fresh generation)")

    # Save prompt (for new/refine modes)
    if generation_mode != "chat":
        (iter_path / "prompt.txt").write_text(prompt)
        logger.info("Saved prompt")

    # ===== REITERATE MODE: Copy prompt from source, skip Claude =====
    if generation_mode == "reiterate":
        import shutil
        source_prompt = char_stage / refined_from / "nanobanana_prompt.txt"
        if not source_prompt.exists():
            return {"status": "error",
                    "error": f"No nanobanana_prompt.txt in {refined_from}"}

        # Copy prompt (will be overwritten with prefix later)
        shutil.copy(source_prompt, iter_path / "nanobanana_prompt.txt")

        # Read prompt, strip any existing [REFERENCE IMAGES...] prefix
        raw_prompt = source_prompt.read_text().strip()
        if raw_prompt.startswith("[REFERENCE IMAGES:"):
            # Strip prefix — it will be re-added with current references
            bracket_end = raw_prompt.find("] ")
            if bracket_end != -1:
                prompt_for_api = raw_prompt[bracket_end + 2:].strip()
            else:
                prompt_for_api = raw_prompt
        else:
            prompt_for_api = raw_prompt

        if "---" in prompt_for_api:
            prompt_for_api = prompt_for_api.split("---")[0].strip()

        is_chat_continuation = False
        chat_history_path = iter_path / "chat_history.json"
        logger.info(f"Reiterate: copied prompt from {refined_from}, skipping Claude")

    # ===== CHAT MODE: Skip Claude, go directly to NanoBanana =====
    elif generation_mode == "chat":
        import shutil

        # Find previous chat history
        prev_chat_history = char_stage / refined_from / "chat_history.json"
        if not prev_chat_history.exists():
            return {"status": "error",
                    "error": f"No chat history in {refined_from}. Use 'refine' instead of 'chat'."}

        # Copy chat history to new iteration
        chat_history_path = iter_path / "chat_history.json"
        shutil.copy(prev_chat_history, chat_history_path)
        logger.info(f"Copied chat history from {refined_from}")

        # The feedback becomes the NanoBanana prompt directly
        prompt_for_api = feedback

        # Save for reference
        (iter_path / "nanobanana_prompt.txt").write_text(
            f"[Chat continuation from {refined_from}]\n\n{feedback}"
        )
        (iter_path / "prompt.txt").write_text(
            f"[CHAT MODE - Direct NanoBanana continuation from {refined_from}]\n\nFeedback: {feedback}"
        )

        # Skip to image generation
        is_chat_continuation = True
        reference_image_path = None  # Chat mode uses history, not reference image

        logger.info("Chat mode: Skipping Claude, using NanoBanana chat continuation")

    else:
        # ===== NEW/REFINE MODES: Call Claude first =====
        is_chat_continuation = False
        chat_history_path = iter_path / "chat_history.json"

    # ===== STEP 1: Call Claude to create NanoBanana prompt (new/refine only) =====
    if generation_mode not in ("chat", "reiterate"):
        logger.info("Step 1: Calling Claude to create NanoBanana prompt...")
        print("\n" + "=" * 60)
        print(f"Stage 6: Characteristic Shot (iteration {iteration})")
        print(f"Session: {session_id}")
        print("=" * 60)

        try:
            result = subprocess.run(
                ["claude", "--resume", session_id,
                 "--allowedTools", "Read,Write,Glob,Grep",
                 "-p", prompt],
                capture_output=True, text=True,
                cwd=str(story_path), timeout=600,
                env=claude_env()
            )

            if result.returncode != 0:
                error_detail = result.stderr or result.stdout
                error_msg = f"Claude returned code {result.returncode}: {error_detail}"
                logger.error(error_msg)
                return {"status": "error", "error": error_msg}

            logger.info("Claude completed successfully")

        except subprocess.TimeoutExpired:
            return {"status": "error", "error": "Claude timed out after 10 minutes"}
        except Exception as e:
            return {"status": "error", "error": f"Exception calling Claude: {str(e)}"}

        # Validate Claude created the prompt file
        nanobanana_prompt_file = iter_path / "nanobanana_prompt.txt"
        if not nanobanana_prompt_file.exists():
            return {"status": "error", "error": "Claude did not create nanobanana_prompt.txt"}

        nanobanana_prompt = nanobanana_prompt_file.read_text().strip()
        logger.info(f"Created nanobanana_prompt.txt ({nanobanana_prompt_file.stat().st_size} bytes)")

        # Extract prompt text (before metadata section)
        prompt_for_api = nanobanana_prompt
        if "---" in prompt_for_api:
            prompt_for_api = prompt_for_api.split("---")[0].strip()

    # ===== STEP 2: Call NanoBanana API to generate image =====
    logger.info("Step 2: Calling NanoBanana API to generate image...")
    print("\n" + "=" * 60)
    print(f"Stage 6: Characteristic Shot (Step 2 - NanoBanana image)")
    if generation_mode == "chat":
        print(f"Mode: CHAT (continuing from {refined_from})")
    print("=" * 60)

    # Load .env for API keys
    try:
        from dotenv import load_dotenv
        env_path = repo_root / ".env"
        if env_path.exists():
            load_dotenv(env_path)
    except ImportError:
        logger.warning("python-dotenv not installed, skipping .env loading")

    output_image = iter_path / "characteristic_shot.png"

    # For NEW/REFINE modes: Use refined_from iteration image as reference
    # (reiterate skips this — same prompt, fresh generation, no extra reference)
    if generation_mode != "chat":
        reference_image_path = None
        ref_roles = []
        if refined_from and generation_mode != "reiterate":
            for ext in ["png", "jpg"]:
                parent_image = char_stage / refined_from / f"characteristic_shot.{ext}"
                if parent_image.exists():
                    reference_image_path = str(parent_image)
                    logger.info(f"Using refined-from iteration as reference: {reference_image_path}")
                    break

        # Initialize chat history path for new generations
        chat_history_path = iter_path / "chat_history.json"
        is_chat_continuation = False
        reference_images = []
        if reference_image_path:
            reference_images.append(reference_image_path)
            ref_roles.append({"role": "refined_from", "source": "iteration"})

        # Add person reference photos from project references
        character_labels = None
        if refs.get("available"):
            if is_story_mode(refs):
                # Story mode: all episode characters + locations + objects
                char_data = get_episode_characters(take_path, refs)
                char_refs = select_character_refs(refs, char_data["all"])
                if char_refs:
                    n_before = len(reference_images)
                    reference_images.extend(r["path"] for r in char_refs)
                    ref_roles.extend(
                        {"role": "character_ref", "asset_id": r["character_id"], "source": "reference"}
                        for r in char_refs
                    )
                    character_labels = [None] * n_before  # None for preceding images
                    character_labels.extend(
                        {"character_id": r["character_id"], "label": r.get("description", "")}
                        for r in char_refs
                    )
                    logger.info(f"Story mode: added {len(char_refs)} character reference sheets")

                # Add location refs with frame index fallback
                from lib.frame_index import load_frame_index
                frame_index = load_frame_index(project_path)

                loc_data = get_episode_locations(take_path, refs)
                loc_refs = select_asset_refs_with_fallback(
                    refs, loc_data["all"], "locations", take_path, frame_index
                )
                remaining = NANOBANANA_MAX_IMAGES - len(reference_images)
                if loc_refs and remaining > 0:
                    loc_refs = loc_refs[:remaining]
                    reference_images.extend(r["path"] for r in loc_refs)
                    ref_roles.extend(
                        {"role": "location_ref", "asset_id": r["asset_id"],
                         "source": "frame" if r["source"] == "frame" else "reference"}
                        for r in loc_refs
                    )
                    if character_labels is not None:
                        character_labels.extend(
                            {"type": "frame_reference", "asset_id": r["asset_id"],
                             "label": r.get("description", "")}
                            if r["source"] == "frame" else None
                            for r in loc_refs
                        )
                    logger.info(f"Story mode: added {len(loc_refs)} location references "
                               f"({sum(1 for r in loc_refs if r['source'] == 'frame')} from frame index)")

                # Add object refs with frame index fallback
                obj_data = get_episode_objects(take_path, refs)
                obj_refs = select_asset_refs_with_fallback(
                    refs, obj_data["all"], "objects", take_path, frame_index
                )
                remaining = NANOBANANA_MAX_IMAGES - len(reference_images)
                if obj_refs and remaining > 0:
                    obj_refs = obj_refs[:remaining]
                    reference_images.extend(r["path"] for r in obj_refs)
                    ref_roles.extend(
                        {"role": "object_ref", "asset_id": r["asset_id"],
                         "source": "frame" if r["source"] == "frame" else "reference"}
                        for r in obj_refs
                    )
                    if character_labels is not None:
                        character_labels.extend(
                            {"type": "frame_reference", "asset_id": r["asset_id"],
                             "label": r.get("description", "")}
                            if r["source"] == "frame" else None
                            for r in obj_refs
                        )
                    logger.info(f"Story mode: added {len(obj_refs)} object references "
                               f"({sum(1 for r in obj_refs if r['source'] == 'frame')} from frame index)")
            else:
                # Product mode: single "main" person
                main_person = refs.get("people", {}).get("main", {})
                person_photos = main_person.get("photos", [])
                if person_photos:
                    photos_to_add = person_photos[:3]
                    reference_images.extend(photos_to_add)
                    ref_roles.extend(
                        {"role": "person_ref", "source": "reference"}
                        for _ in photos_to_add
                    )
                    logger.info(f"Added {len(photos_to_add)} person reference photos")

        reference_images = reference_images or None

    else:
        # CHAT mode: reference_images is None, history is already set up
        reference_images = None
        ref_roles = []
        character_labels = None

    # Prepend [REFERENCE IMAGES...] prefix for reference image guidance + sharpness
    if reference_images and generation_mode != "chat":
        refs_path = project_path / "references" if refs.get("available") else None
        index = refs.get("index")
        analyze_prefix = build_analyze_prefix(
            index, reference_images, refs_path,
            character_labels=character_labels
        )
        if analyze_prefix:
            prompt_for_api = analyze_prefix + prompt_for_api
            # Save the full prompt (with reference prefix) for auditability
            (iter_path / "nanobanana_prompt.txt").write_text(prompt_for_api)
            logger.info(f"Added [REFERENCE IMAGES...] prefix ({len(analyze_prefix)} chars)")

    # Scale down references if requested
    if scale_refs and reference_images:
        from lib import scale_down_references
        scaled_dir = iter_path / "scaled_refs"
        reference_images = scale_down_references(
            reference_images, scaled_dir, max_pixels=2048
        )
        logger.info(f"Scaled references to 2K in {scaled_dir}")

    try:
        from nanobanana_text2img import generate_image_chat

        result_path = generate_image_chat(
            prompt=prompt_for_api,
            output_path=str(output_image),
            history_path=str(chat_history_path),
            aspect_ratio="16:9",
            image_size="4K",
            reference_image_paths=reference_images,
            is_refinement=is_chat_continuation
        )

        logger.info(f"Image generated: {result_path}")
        output_image = Path(result_path)

    except Exception as e:
        error_msg = f"NanoBanana image generation failed: {str(e)}"
        logger.error(error_msg)
        return {"status": "error", "error": error_msg}

    # User message
    user_message_file = iter_path / "user_message.txt"
    user_message = (user_message_file.read_text().strip()
                    if user_message_file.exists()
                    else "Characteristic shot generated. Please review the image.")

    # Save metadata (V5 format)
    metadata = {
        "stage": "06-characteristic",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"06-characteristic/{iteration}/{output_image.name}",
        "generation_mode": generation_mode,
        "model_prompt": "direct-chat" if generation_mode == "chat" else ("reused" if generation_mode == "reiterate" else "claude-default"),
        "model_image": "nanobanana-gemini3pro",
        "reference_images_sent": build_reference_record(
            reference_images or [], project_path, ref_roles
        ),
        "session_id": session_id,
        "image_file": output_image.name,
        "aspect_ratio": "16:9",
        "chat_history": "chat_history.json" if chat_history_path.exists() else None,
    }

    save_metadata(iter_path / "metadata.json", metadata)
    logger.info("Saved metadata")

    # Log event
    log_event(
        story_path, "stage_complete",
        {"action": "characteristic_generated", "iteration": iteration,
         "result": f"06-characteristic/{iteration}/{output_image.name}"},
        take=take_path.name, stage="06-characteristic", iteration=iteration
    )

    print(f"\nGenerated: 06-characteristic/{iteration}/{output_image.name}")

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"06-characteristic/{iteration}/{output_image.name}",
        "user_message": user_message,
    }
