#!/usr/bin/env python3
"""
Stage S3: Visual Style + Style Board

Claude defines visual style. NanoBanana generates style board image.
Title artwork is handled separately in S4.
"""

import subprocess
import sys
import os
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
from story_stages.stage_s2_bible import _build_references_section


def execute(story_path: Path, inputs: dict,
            feedback: str = None, refined_from: str = None,
            style_preference: str = None) -> dict:
    """
    Execute Stage S3: Style + Style Board.

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

    print(f"\nStage S3: Style + Style Board (iteration {iteration})")

    # Resolve bible path
    bible_iter = inputs.get("S2-bible")
    if not bible_iter:
        return {"status": "error", "error": "S2-bible not in inputs"}
    bible_path = f"S2-bible/{bible_iter}/story_bible.json"
    bible_full_path = story_path / bible_path
    if not bible_full_path.exists():
        return {"status": "error", "error": f"Story bible not found: {bible_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_s3_style.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)

    # Style preference
    if style_preference:
        prompt = prompt.replace("{user_style_preference}",
                                f"## User Style Preference\n\n{style_preference}")
    else:
        prompt = prompt.replace("{user_style_preference}", "")

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

    # Add feedback
    if refined_from:
        fb_history = get_feedback_history_for_claude(story_path, "S3-style", 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 for style definition + NanoBanana prompt
    print(f"Session: {session_id}")
    print("Calling Claude to define visual style...")

    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"}

    # Validate Claude outputs
    style_file = iter_path / "style.json"
    if not style_file.exists():
        return {"status": "error", "error": "style.json not created"}

    # Step 2: Generate style board via NanoBanana
    style_board_prompt_file = iter_path / "nanobanana_style_board.txt"

    if style_board_prompt_file.exists():
        print("Generating style board via NanoBanana...")
        style_board_result = _generate_nanobanana_image(
            prompt_file=style_board_prompt_file,
            output_path=iter_path / "style_board.jpg",
            chat_history_path=iter_path / "chat_history.json",
        )
        if style_board_result:
            # Copy to references
            refs_style_dir = story_path / "references" / "style"
            ensure_dir(refs_style_dir)
            shutil.copy2(style_board_result, refs_style_dir / "style_board.jpg")
            print(f"  → references/style/style_board.jpg")
        else:
            print("  Warning: Style board generation failed")
    else:
        print("  Warning: No NanoBanana prompt for style board")

    # User message
    user_msg_file = iter_path / "user_message.txt"
    user_message = (user_msg_file.read_text().strip() if user_msg_file.exists()
                    else "Visual style defined. Review style.json and style board.")

    # Save metadata
    has_style_board = (iter_path / "style_board.jpg").exists()
    metadata = {
        "stage": "S3-style",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"S3-style/{iteration}/style.json",
        "session_id": session_id,
        "has_style_board": has_style_board,
        "style_preference": style_preference,
    }
    save_metadata(iter_path / "metadata.json", metadata)

    print(f"\nGenerated: S3-style/{iteration}/")
    print(f"  style.json" + (" + style_board.jpg" if has_style_board else ""))

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"S3-style/{iteration}/style.json",
        "user_message": user_message,
    }


def _generate_nanobanana_image(prompt_file: Path, output_path: Path,
                                chat_history_path: Path) -> Path | None:
    """Generate an image via NanoBanana API.

    Returns output path on success, None on failure.
    """
    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",
        )
        print(f"  Generated: {Path(result_path).name}")
        return Path(result_path)

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