#!/usr/bin/env python3
"""
Stage S5: Character & Location Reference Design

Claude creates NanoBanana prompts for each character portrait and key location.
NanoBanana generates reference images.
Results are copied to references/people/{id}/ and references/locations/.
"""

import subprocess

from lib import claude_env
import sys
import json
import shutil
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
from story_stages.stage_s2_bible import _build_references_section


def execute(story_path: Path, inputs: dict,
            feedback: str = None, refined_from: str = None) -> dict:
    """
    Execute Stage S4: Character & Location Design.

    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 / "S5-characters"
    ensure_dir(stage_dir)
    iteration = get_next_iteration(stage_dir)
    iter_path = stage_dir / iteration
    ensure_dir(iter_path)

    print(f"\nStage S5: Character & Location Design (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}"}

    # 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_s5_characters.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)

    refs_section = _build_references_section(story_path)
    prompt = prompt.replace("{user_references_section}", refs_section)

    if refined_from:
        fb_history = get_feedback_history_for_claude(story_path, "S5-characters", 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 prompts + descriptions
    print(f"Session: {session_id}")
    print("Calling Claude to design characters and locations...")

    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: Read characters_plan.json and generate images
    plan_file = iter_path / "characters_plan.json"
    if not plan_file.exists():
        return {"status": "error", "error": "characters_plan.json not created by Claude"}

    plan = json.loads(plan_file.read_text())

    images_generated = []
    refs_path = story_path / "references"

    # Generate character portraits
    for char_info in plan.get("characters", []):
        char_id = char_info["id"]
        prompt_file = iter_path / char_info.get("prompt_file", f"{char_id}_prompt.txt")
        desc_file = iter_path / char_info.get("description_file", f"{char_id}_description.md")

        if not prompt_file.exists():
            print(f"  Warning: No prompt file for {char_id}")
            continue

        print(f"Generating portrait for {char_id}...")

        # Create character directory in iteration
        char_dir = iter_path / char_id
        ensure_dir(char_dir)

        output_path = char_dir / "portrait.jpg"
        chat_history = char_dir / "chat_history.json"

        result_path = _generate_nanobanana_image(
            prompt_file=prompt_file,
            output_path=output_path,
            chat_history_path=chat_history,
        )

        if result_path:
            images_generated.append(f"{char_id}/portrait.jpg")

            # Copy to references/people/{id}/
            ref_char_dir = refs_path / "people" / char_id
            ensure_dir(ref_char_dir)
            shutil.copy2(result_path, ref_char_dir / "portrait.jpg")

            # Copy description to references
            if desc_file.exists():
                shutil.copy2(desc_file, ref_char_dir / "description.md")

            print(f"  → references/people/{char_id}/portrait.jpg")
        else:
            print(f"  Warning: Portrait generation failed for {char_id}")

    # Generate location references
    for loc_info in plan.get("locations", []):
        loc_id = loc_info["id"]
        prompt_file = iter_path / loc_info.get("prompt_file", f"{loc_id}_prompt.txt")

        if not prompt_file.exists():
            print(f"  Warning: No prompt file for location {loc_id}")
            continue

        print(f"Generating location reference for {loc_id}...")

        output_path = iter_path / f"{loc_id}.jpg"
        chat_history = iter_path / f"{loc_id}_chat_history.json"

        result_path = _generate_nanobanana_image(
            prompt_file=prompt_file,
            output_path=output_path,
            chat_history_path=chat_history,
        )

        if result_path:
            images_generated.append(f"{loc_id}.jpg")

            # Copy to references/locations/
            ref_loc_dir = refs_path / "locations"
            ensure_dir(ref_loc_dir)
            shutil.copy2(result_path, ref_loc_dir / f"{loc_id}.jpg")
            print(f"  → references/locations/{loc_id}.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 f"Generated {len(images_generated)} reference images.")

    # Save metadata
    metadata = {
        "stage": "S5-characters",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"S5-characters/{iteration}/",
        "session_id": session_id,
        "images_generated": images_generated,
        "characters": [c["id"] for c in plan.get("characters", [])],
        "locations": [l["id"] for l in plan.get("locations", [])],
    }
    save_metadata(iter_path / "metadata.json", metadata)

    print(f"\nGenerated: S5-characters/{iteration}/")
    print(f"  {len(images_generated)} images total")

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"S5-characters/{iteration}/",
        "user_message": user_message,
        "images": images_generated,
    }


def _generate_nanobanana_image(prompt_file: Path, output_path: Path,
                                chat_history_path: Path) -> 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:
        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",
        )
        return Path(result_path)

    except Exception as e:
        print(f"  NanoBanana error: {e}")
        return None
