"""
Stage 11: Title Card Generation

Generates intro and outro title card images using NanoBanana.
Scene-based stage with 2 fixed scenes: "intro" and "outro".

This stage runs AFTER video generation (stage 10) so title content
can reflect actual episode content.

Commands:
- generate_title: Generate title card (intro or outro)
- refine_title: Refine title card with NanoBanana chat
- select_title: Select title card iteration
- get_titles_status: Show status for both intro and outro
"""
import json
import subprocess
import sys
from pathlib import Path
from datetime import datetime, timezone
from typing import Optional, List

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

from lib.path_utils_v3 import ensure_dir
from lib.generation_log import setup_generation_log
from lib.pipeline import load_pipeline, save_pipeline
from lib.image_utils import scale_down_references
from lib import claude_env
from nanobanana_text2img import generate_image, generate_image_chat


def generate_title(
    take_path: Path,
    title_type: str,
    feedback: Optional[str] = None
) -> dict:
    """
    Generate a title card image (intro or outro).

    Uses Claude to design the NanoBanana prompt based on episode content
    (adapted script + style + characteristic shot), then NanoBanana generates
    the image.

    Args:
        take_path: Path to episode/take directory
        title_type: "intro" or "outro"
        feedback: Optional user guidance for prompt generation

    Returns:
        Result dict with status, frame name
    """
    if title_type not in ("intro", "outro"):
        return {"status": "error", "error": f"Invalid title_type: {title_type}. Must be 'intro' or 'outro'."}

    logger = setup_generation_log(take_path)
    logger.section(f"Stage 11: Generate Title Card - {title_type}")

    # Create title directory
    titles_dir = take_path / "11-titles"
    title_dir = titles_dir / title_type
    ensure_dir(title_dir)

    # Determine iteration number
    existing = sorted([
        p for p in title_dir.iterdir()
        if p.suffix == ".jpg" and p.stem.isdigit()
    ]) if title_dir.exists() else []
    next_num = int(existing[-1].stem) + 1 if existing else 1
    frame_name = f"{next_num:03d}"
    output_path = title_dir / f"{frame_name}.jpg"

    # Collect reference images BEFORE Claude call so we can tell Claude what's attached
    ref_images, ref_labels = _collect_title_references(take_path, title_type)
    if ref_images:
        scaled_dir = title_dir / "scaled_refs"
        ref_images = scale_down_references(ref_images, scaled_dir)
        logger.info(f"Reference images: {len(ref_images)}")

    # Generate prompt using Claude (pass ref_labels so Claude writes [REFERENCE IMAGES:] block)
    prompt = _generate_title_prompt(take_path, title_type, feedback, logger, ref_labels)
    if not prompt:
        return {"status": "error", "error": "Failed to generate title card prompt"}

    logger.info(f"Generating {title_type} title card -> {frame_name}.jpg")
    logger.info(f"Prompt: {prompt[:100]}...")

    # Save prompt
    prompt_file = title_dir / f"prompt-{frame_name}.txt"
    prompt_file.write_text(prompt)

    try:
        result_path = generate_image(
            prompt=prompt,
            output_path=str(output_path),
            aspect_ratio="9:16",
            image_size="4K",
            reference_image_path=ref_images[0] if ref_images else None,
            additional_reference_paths=ref_images[1:] if len(ref_images) > 1 else None
        )

        if not result_path or not Path(result_path).exists():
            logger.error("NanoBanana title card generation failed")
            return {"status": "error", "error": "Image generation failed"}

        logger.info(f"Created {title_type}/{frame_name}.jpg")

        # Update metadata
        meta_file = title_dir / "metadata.json"
        if meta_file.exists():
            with open(meta_file) as f:
                meta = json.load(f)
        else:
            meta = {"frames": {}, "selected": None, "selected_at": None}

        meta["frames"][frame_name] = {
            "type": "generated",
            "title_type": title_type,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "rating": None
        }

        with open(meta_file, 'w') as f:
            json.dump(meta, f, indent=2)

        logger.info(f"Next: Review 11-titles/{title_type}/{frame_name}.jpg")

        return {"status": "success", "frame": frame_name, "title_type": title_type}

    except Exception as e:
        logger.error(f"Title card generation failed: {e}")
        import traceback
        traceback.print_exc()
        return {"status": "error", "error": str(e)}


def refine_title(
    take_path: Path,
    title_type: str,
    base_frame: str,
    feedback: str
) -> dict:
    """
    Refine title card using NanoBanana chat mode.

    Args:
        take_path: Path to episode directory
        title_type: "intro" or "outro"
        base_frame: Base frame number (e.g., "001")
        feedback: User feedback for refinement

    Returns:
        Result dict with status, new frame name
    """
    if title_type not in ("intro", "outro"):
        return {"status": "error", "error": f"Invalid title_type: {title_type}"}

    logger = setup_generation_log(take_path)
    logger.section(f"Stage 11: Refine Title Card - {title_type}/{base_frame}")

    titles_dir = take_path / "11-titles"
    title_dir = titles_dir / title_type

    if not title_dir.exists():
        return {"status": "error",
                "error": f"Title card ({title_type}) not generated yet. Run generate-title first."}

    base_frame_path = title_dir / f"{base_frame}.jpg"
    if not base_frame_path.exists():
        logger.error(f"Base frame not found: {base_frame}.jpg")
        return {"status": "error", "error": f"Frame not found: {base_frame}.jpg"}

    # Chat history path
    chat_history_path = title_dir / f"nb-chat-history-{base_frame}.json"
    is_first_refinement = not chat_history_path.exists()

    # Count existing refinements
    existing_refinements = []
    for item in title_dir.iterdir():
        if item.name.startswith(f"{base_frame}-refined-") and item.suffix == ".jpg":
            try:
                iter_num = int(item.name.split('-')[-1].replace('.jpg', ''))
                existing_refinements.append(iter_num)
            except ValueError:
                pass

    next_iter = max(existing_refinements) + 1 if existing_refinements else 1
    output_frame_name = f"{base_frame}-refined-{next_iter:03d}"
    output_frame_path = title_dir / f"{output_frame_name}.jpg"

    logger.info(f"Refining {title_type} title card {base_frame}.jpg -> {output_frame_name}.jpg")
    logger.info(f"Feedback: {feedback}")

    # Build prompt
    if is_first_refinement:
        prompt = (
            f"This is a {title_type} title card for a video series. "
            f"Refine based on this feedback:\n\n{feedback}\n\n"
            f"TECHNICAL REQUIREMENTS:\n"
            f"- Maintain 9:16 aspect ratio\n"
            f"- High sharpness and detail\n"
            f"- Keep text legible and well-composed"
        )
    else:
        # Chat continuation
        prompt = feedback

    # Save prompt
    prompt_file = title_dir / f"prompt-{output_frame_name}.txt"
    prompt_file.write_text(prompt)

    try:
        ref_images = []
        if is_first_refinement:
            ref_images.append(str(base_frame_path))

        result_path = generate_image_chat(
            prompt=prompt,
            output_path=str(output_frame_path),
            history_path=str(chat_history_path),
            aspect_ratio="9:16",
            image_size="4K",
            reference_image_paths=ref_images if ref_images else None,
            is_refinement=(not is_first_refinement)
        )

        if not result_path or not Path(result_path).exists():
            logger.error("NanoBanana title card refinement failed")
            return {"status": "error", "error": "Image generation failed"}

        logger.info(f"Created {title_type}/{output_frame_name}.jpg")

        # Update metadata
        meta_file = title_dir / "metadata.json"
        with open(meta_file) as f:
            meta = json.load(f)

        meta["frames"][output_frame_name] = {
            "type": "refined",
            "title_type": title_type,
            "base_frame": f"{base_frame}.jpg",
            "chat_turn": next_iter + 1,
            "feedback": feedback,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "rating": None
        }

        with open(meta_file, 'w') as f:
            json.dump(meta, f, indent=2)

        # Save feedback log
        feedback_log = title_dir / f"feedback-{base_frame}.txt"
        with open(feedback_log, 'a') as f:
            f.write(f"[{next_iter:03d}] {feedback}\n")

        return {
            "status": "success",
            "frame": output_frame_name,
            "title_type": title_type,
            "chat_turn": next_iter + 1
        }

    except Exception as e:
        logger.error(f"Title card refinement failed: {e}")
        import traceback
        traceback.print_exc()
        return {"status": "error", "error": str(e)}


def select_title(take_path: Path, title_type: str, frame: str) -> dict:
    """
    Select a title card frame for intro or outro.

    Updates pipeline["11-titles"][title_type] = frame.

    Args:
        take_path: Path to take directory
        title_type: "intro" or "outro"
        frame: Frame name (e.g., "001" or "001-refined-002")

    Returns:
        Result dict with status
    """
    if title_type not in ("intro", "outro"):
        return {"status": "error", "error": f"Invalid title_type: {title_type}"}

    logger = setup_generation_log(take_path)
    logger.section(f"Stage 11: Select Title - {title_type}")

    titles_dir = take_path / "11-titles"
    title_dir = titles_dir / title_type

    # Validate frame exists
    frame_path = title_dir / f"{frame}.jpg"
    if not frame_path.exists():
        return {"status": "error", "error": f"Frame not found: {frame}.jpg"}

    # Update main pipeline
    pipeline = load_pipeline(take_path)

    if not isinstance(pipeline.get("11-titles"), dict):
        pipeline["11-titles"] = {}

    pipeline["11-titles"][title_type] = frame
    save_pipeline(take_path, pipeline)

    # Update title metadata
    meta_file = title_dir / "metadata.json"
    if meta_file.exists():
        with open(meta_file) as f:
            meta = json.load(f)
        meta["selected"] = frame
        meta["selected_at"] = datetime.now(timezone.utc).isoformat()
        with open(meta_file, 'w') as f:
            json.dump(meta, f, indent=2)

    logger.info(f"Selected {frame}.jpg for {title_type} title")
    return {"status": "success", "title_type": title_type, "frame": frame}


def get_titles_status(take_path: Path) -> dict:
    """
    Get status for both intro and outro title cards.

    Returns:
        {
            "status": "success",
            "intro": {"generated": bool, "selected": str|None, "available": [...]},
            "outro": {"generated": bool, "selected": str|None, "available": [...]}
        }
    """
    pipeline = load_pipeline(take_path)
    title_selections = pipeline.get("11-titles")
    if not isinstance(title_selections, dict):
        title_selections = {}

    result = {"status": "success"}

    for title_type in ("intro", "outro"):
        title_dir = take_path / "11-titles" / title_type

        if not title_dir.exists():
            result[title_type] = {"generated": False, "selected": None, "available": [], "count": 0}
            continue

        available = sorted([
            p.stem for p in title_dir.iterdir()
            if p.suffix == ".jpg"
        ])

        result[title_type] = {
            "generated": len(available) > 0,
            "selected": title_selections.get(title_type),
            "available": available,
            "count": len(available)
        }

    return result


def _generate_title_prompt(
    take_path: Path,
    title_type: str,
    feedback: Optional[str],
    logger,
    ref_labels: Optional[List[str]] = None
) -> Optional[str]:
    """
    Use Claude to generate a NanoBanana prompt for the title card.

    Analyzes adapted script + characteristic shot + video metadata
    to design an appropriate title card. Includes reference image
    descriptions so Claude can write the [REFERENCE IMAGES:] block.
    """
    # Generate prompt using Claude
    session_file = take_path.parent / "session.txt"
    if not session_file.exists():
        logger.warning("No session file found, cannot generate title prompt")
        return None

    session_id = session_file.read_text().strip()
    story_path = take_path.parent

    # Gather context
    pipeline = load_pipeline(take_path)

    # Get style info from adapted script
    style_info = ""
    script_iter = pipeline.get("09-script-adapted")
    if script_iter:
        script_file = take_path / "09-script-adapted" / script_iter / "script.json"
        if script_file.exists():
            try:
                script = json.loads(script_file.read_text())
                title = script.get("title", "")
                style = script.get("style", {})
                visual = style.get("visual", "")
                reference = style.get("reference", "")
                style_info = f"Title: {title}\nVisual style: {visual}\nReference: {reference}"
            except Exception:
                pass

    # Get episode info
    episode_info = take_path.name

    # Get channel name from channel_concept.txt (first line, after @)
    # Falls back to directory name if not available
    story_name = story_path.name
    concept_file = story_path / "channel_concept.txt"
    if concept_file.exists():
        first_line = concept_file.read_text().strip().split('\n')[0]
        # Extract @channel.name from first line like "# @no.fix.planned — ..."
        import re
        match = re.search(r'@[\w.]+', first_line)
        if match:
            story_name = match.group(0)

    # Load reference prompt from references/titles/ if available
    ref_prompt_text = ""
    refs_titles_dir = story_path / "references" / "titles"
    ref_prompt_file = refs_titles_dir / f"{title_type}_prompt_reference.txt"
    if ref_prompt_file.exists():
        ref_prompt_text = ref_prompt_file.read_text().strip()
        logger.info(f"Loaded reference prompt: {ref_prompt_file.name}")

    # Build reference images description for Claude
    ref_images_section = ""
    if ref_labels:
        image_list = "\n".join(f"  Image {i}: {label}" for i, label in enumerate(ref_labels, 1))
        ref_images_section = f"""
REFERENCE IMAGES THAT WILL BE ATTACHED TO THE PROMPT:
{image_list}

You MUST start your prompt with a [REFERENCE IMAGES: ...] block that describes each image's role.
Image 1 is the PRIMARY STYLE REFERENCE (previous episode's title card) — the generated image
must match its composition, lettering style, and color palette.
All other images are CONTENT REFERENCES — use them to ensure characters, objects, and locations
look correct. Do NOT copy their composition.
"""

    if title_type == "intro":
        ref_prompt_section = ""
        if ref_prompt_text:
            ref_prompt_section = f"""
REFERENCE PROMPT (from Episode 1 — match this style and structure):
---
{ref_prompt_text}
---

Generate a similar prompt for THIS episode. Keep the same lettering style,
color palette, vintage sci-fi paperback feel. Change only the episode-specific
content (episode number, subtitle, background scene).
"""

        claude_prompt = f"""Design a title card image for the INTRO of a short video episode.

EPISODE INFO:
Story: {story_name}
Episode: {episode_info}
{style_info}
{ref_prompt_section}{ref_images_section}
REQUIREMENTS:
- This is the OPENING title card shown before the episode starts
- Include the series title and episode subtitle, artistically integrated
- The title text should be LEGIBLE but stylish — not just plain white text
- Background should set the mood for the episode
- 9:16 vertical format (1080x1920)
- Keep the same overall structure as the reference: series title at top, background scene, episode info at bottom

{f"USER GUIDANCE: {feedback}" if feedback else ""}

Write ONLY the NanoBanana image generation prompt (starting with [REFERENCE IMAGES:] block).
No explanations, no markdown.
"""
    else:  # outro
        ref_prompt_section = ""
        if ref_prompt_text:
            ref_prompt_section = f"""
REFERENCE PROMPT (from Episode 1 — match this style and structure):
---
{ref_prompt_text}
---

Generate a similar prompt for THIS episode. Keep the same visual style,
typography, and composition. Change only the episode-specific text.
"""

        # Load title_card_notes from project config
        title_card_notes = ""
        sp_file = story_path / "story_pipeline.json"
        if sp_file.exists():
            try:
                sp = json.loads(sp_file.read_text())
                notes = sp.get("title_card_notes")
                if notes:
                    title_card_notes = f"\nPROJECT-SPECIFIC NOTES:\n{notes}\n"
            except (json.JSONDecodeError, OSError):
                pass

        claude_prompt = f"""Design a title card image for the OUTRO of a short video episode.

EPISODE INFO:
Story: {story_name}
Episode: {episode_info}
{style_info}
{ref_prompt_section}{ref_images_section}{title_card_notes}
REQUIREMENTS:
- This is the CLOSING card shown after the episode ends
- Include "To be continued..." or similar text, artistically styled
- Include series name and episode info
- Should create anticipation for the next episode
- 9:16 vertical format (1080x1920)

{f"USER GUIDANCE: {feedback}" if feedback else ""}

Write ONLY the NanoBanana image generation prompt (starting with [REFERENCE IMAGES:] block).
No explanations, no markdown.
"""

    try:
        result = subprocess.run(
            ["claude", "--resume", session_id,
             "--model", "opus",
             "-p", claude_prompt],
            capture_output=True,
            text=True,
            cwd=str(story_path),
            timeout=120,
            env=claude_env()
        )

        if result.returncode == 0 and result.stdout.strip():
            return result.stdout.strip()
        else:
            logger.warning(f"Claude prompt generation failed: {result.stderr or result.stdout}")
            return None

    except Exception as e:
        logger.warning(f"Claude error: {e}")
        return None


def _collect_title_references(
    take_path: Path,
    title_type: str = "intro"
) -> tuple:
    """Collect reference images + labels for title card generation.

    Uses assets.json to find episode-relevant characters and objects.
    Returns (ref_paths, ref_labels) where:
    - Image 1: previous episode's title card (primary style reference)
    - Images 2+: episode-relevant character sheets and object refs

    Title refs ensure consistent series look; character/object refs ensure
    the title card depicts the RIGHT characters and objects.
    """
    refs = []
    labels = []

    project_path = take_path.parent
    refs_path = project_path / "references"

    # 1. Primary: previous episode's title card (style reference)
    titles_dir = refs_path / "titles"
    if titles_dir.exists():
        for img in sorted(titles_dir.iterdir()):
            if img.suffix in (".jpg", ".png") and title_type in img.stem:
                refs.append(str(img))
                labels.append(f"previous episode {title_type} title card (PRIMARY STYLE REFERENCE — match composition, lettering, colors)")
                break

    # Load pipeline + assets.json to find episode-relevant characters/objects
    pipeline = load_pipeline(take_path)

    # 2. Characteristic shot — overall visual style of this episode
    char_iter = pipeline.get("06-characteristic")
    if char_iter:
        char_shot = take_path / "06-characteristic" / char_iter / "characteristic_shot.jpg"
        if char_shot.exists():
            refs.append(str(char_shot))
            labels.append("episode characteristic shot (overall visual style, composition, color palette)")

    ep_chars = set()
    ep_objects = set()
    assets_iter = pipeline.get("04-assets")
    if assets_iter:
        assets_file = take_path / "04-assets" / assets_iter / "assets.json"
        if assets_file.exists():
            try:
                assets = json.loads(assets_file.read_text())
                for char in assets.get("characters", []):
                    if char.get("existing_references"):
                        ep_chars.add(char["id"])
                for obj in assets.get("objects", []):
                    if obj.get("existing_references"):
                        ep_objects.add(obj["id"])
            except Exception:
                pass

    # 3. Character sheets — only episode-relevant characters
    people_path = refs_path / "people"
    if people_path.exists():
        for char_dir in sorted(people_path.iterdir()):
            if not char_dir.is_dir():
                continue
            # Skip characters not in this episode (if we have assets data)
            if ep_chars and char_dir.name not in ep_chars:
                continue
            char_label = char_dir.name.replace("_", " ").title()
            for img_name in ["sheet.jpg", "portrait.jpg"]:
                img_path = char_dir / img_name
                if img_path.exists():
                    refs.append(str(img_path))
                    kind = "character sheet (face/body/clothing reference)" if "sheet" in img_name else "portrait (face reference)"
                    labels.append(f"{char_label} {kind}")
                    break

    # 4. Object references — only episode-relevant objects
    # Read channel name for logo label
    channel_name = None
    concept_file = project_path / "channel_concept.txt"
    if concept_file.exists():
        import re as _re
        first_line = concept_file.read_text().strip().split('\n')[0]
        m = _re.search(r'@[\w.]+', first_line)
        if m:
            channel_name = m.group(0)

    objects_path = refs_path / "objects"
    if objects_path.exists():
        for obj_item in sorted(objects_path.iterdir()):
            if obj_item.suffix in (".jpg", ".png"):
                obj_id = obj_item.stem
                if ep_objects and obj_id not in ep_objects:
                    continue
                obj_label = obj_id.replace("_", " ").title()
                refs.append(str(obj_item))
                # For channel_logo, include actual channel name so Claude doesn't guess
                if obj_id == "channel_logo" and channel_name:
                    labels.append(f"channel logo for \"{channel_name}\" — reproduce this logo EXACTLY as shown")
                else:
                    labels.append(f"{obj_label} object reference")

    # 5. Title artwork (S4-title) — lettering/logo style
    if titles_dir and titles_dir.exists():
        artwork = titles_dir / "title_artwork.jpg"
        if artwork.exists():
            refs.append(str(artwork))
            labels.append("series title artwork (lettering/logo reference)")

    return refs, labels
