#!/usr/bin/env python3
"""
Stage S4: Title Artwork

Generates series title artwork via NanoBanana, using the style board
from S3 as a reference image for visual consistency.
"""

import subprocess
import sys
import shutil

from lib import claude_env
from pathlib import Path
from datetime import datetime, timezone

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
from lib.path_utils_v3 import get_next_iteration
from lib.feedback_history_v5 import get_feedback_history_for_claude


def execute(story_path: Path, inputs: dict,
            feedback: str = None, refined_from: str = None) -> dict:
    """
    Execute Stage S4: Title Artwork.

    Args:
        story_path: Path to story project directory
        inputs: {"S2-bible": "001", "S3-style": "001"}
        feedback: User feedback
        refined_from: Iteration being refined
    """
    stage_dir = story_path / "S4-title"
    ensure_dir(stage_dir)
    iteration = get_next_iteration(stage_dir)
    iter_path = stage_dir / iteration
    ensure_dir(iter_path)

    print(f"\nStage S4: Title Artwork (iteration {iteration})")

    # Resolve inputs
    bible_iter = inputs.get("S2-bible")
    style_iter = inputs.get("S3-style")
    if not bible_iter:
        return {"status": "error", "error": "S2-bible not in inputs"}
    if not style_iter:
        return {"status": "error", "error": "S3-style not in inputs"}

    bible_path = f"S2-bible/{bible_iter}/story_bible.json"
    style_path = f"S3-style/{style_iter}/style.json"

    if not (story_path / bible_path).exists():
        return {"status": "error", "error": f"Story bible not found: {bible_path}"}
    if not (story_path / style_path).exists():
        return {"status": "error", "error": f"Style not found: {style_path}"}

    # Find style board as reference image
    style_board_path = story_path / "references" / "style" / "style_board.jpg"
    if not style_board_path.exists():
        # Try in the style iteration directory
        style_board_path = story_path / f"S3-style/{style_iter}/style_board.jpg"

    # Load session
    session_file = story_path / "session.txt"
    if not session_file.exists():
        return {"status": "error", "error": "session.txt not found"}
    session_id = session_file.read_text().strip()

    # Build prompt
    prompt_path = Path(__file__).parent.parent.parent / "prompts" / "story" / "stage_s4_title.txt"
    if not prompt_path.exists():
        return {"status": "error", "error": f"Prompt template not found: {prompt_path}"}

    prompt = prompt_path.read_text()
    prompt = prompt.replace("{bible_path}", bible_path)
    prompt = prompt.replace("{style_path}", style_path)

    # Add feedback
    if refined_from:
        fb_history = get_feedback_history_for_claude(story_path, "S4-title", refined_from)
        if fb_history:
            prompt += f"\n\n## Previous Feedback\n\n{fb_history}\n"
    if feedback:
        prompt += f"\n\n## User Feedback\n\n{feedback}\n"
        (iter_path / "feedback.txt").write_text(feedback)

    # Output directory
    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}"

    (iter_path / "prompt.txt").write_text(prompt)

    # Step 1: Call Claude to create NanoBanana prompt for title artwork
    print(f"Session: {session_id}")
    print("Calling Claude to design title artwork...")

    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
            return {"status": "error",
                    "error": f"Claude returned code {result.returncode}: {error_detail}"}
    except subprocess.TimeoutExpired:
        return {"status": "error", "error": "Claude timed out"}

    # Step 2: Generate title artwork via NanoBanana with style board as reference
    title_prompt_file = iter_path / "nanobanana_title_artwork.txt"

    if not title_prompt_file.exists():
        return {"status": "error", "error": "nanobanana_title_artwork.txt not created by Claude"}

    print("Generating title artwork via NanoBanana...")

    # Use style board as reference for visual consistency
    ref_images = []
    if style_board_path.exists():
        ref_images = [str(style_board_path)]
        print(f"  Using style board as reference: {style_board_path.name}")

    title_result = _generate_nanobanana_image(
        prompt_file=title_prompt_file,
        output_path=iter_path / "title_artwork.jpg",
        chat_history_path=iter_path / "chat_history.json",
        reference_images=ref_images,
    )

    if not title_result:
        return {"status": "error", "error": "Title artwork generation failed"}

    # Copy to references
    refs_style_dir = story_path / "references" / "style"
    ensure_dir(refs_style_dir)
    shutil.copy2(title_result, refs_style_dir / "title_artwork.jpg")
    print(f"  → references/style/title_artwork.jpg")

    # User message
    user_msg_file = iter_path / "user_message.txt"
    user_message = (user_msg_file.read_text().strip() if user_msg_file.exists()
                    else "Title artwork generated.")

    # Save metadata
    metadata = {
        "stage": "S4-title",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"S4-title/{iteration}/title_artwork.jpg",
        "session_id": session_id,
        "has_reference": bool(ref_images),
    }
    save_metadata(iter_path / "metadata.json", metadata)

    print(f"\nGenerated: S4-title/{iteration}/title_artwork.jpg")

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"S4-title/{iteration}/title_artwork.jpg",
        "user_message": user_message,
    }


def _generate_nanobanana_image(prompt_file: Path, output_path: Path,
                                chat_history_path: Path,
                                reference_images: list = None) -> Path | None:
    """Generate an image via NanoBanana API."""
    try:
        from dotenv import load_dotenv
        env_path = Path(__file__).parent.parent.parent.parent / ".env"
        if env_path.exists():
            load_dotenv(env_path)
    except ImportError:
        pass

    prompt_text = prompt_file.read_text().strip()
    if not prompt_text:
        print(f"  Warning: Empty prompt in {prompt_file.name}")
        return None

    try:
        from nanobanana_text2img import generate_image_chat

        result_path = generate_image_chat(
            prompt=prompt_text,
            output_path=str(output_path),
            history_path=str(chat_history_path),
            aspect_ratio="9:16",
            image_size="4K",
            reference_image_paths=reference_images,
        )
        print(f"  Generated: {Path(result_path).name}")
        return Path(result_path)

    except Exception as e:
        print(f"  Warning: NanoBanana generation failed: {e}")
        return None
