#!/usr/bin/env python3
"""
Stage 9: Script Adaptation (V5 orchestrator)

Reads original script from 03-script/{inputs["03-script"]}/script.json
and selected frames from 08-frames/ (per-scene selections + refinement feedback).

Uses Claude with vision to analyze the actual selected/refined frames and
adapt the script to match what was actually generated.

Creates 09-script-adapted/XXX/script.json.

V5 pattern: uses inputs dict for dependencies, get_next_iteration, no symlinks.

NOTE: This stage now runs AFTER frame selection (stage 08), not before it.
This ensures the adapted script describes the actual frames the user selected
and refined, rather than a grid that may not match final frame choices.
"""

import subprocess
import sys
import json
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))

from lib import (
    save_metadata,
    ensure_dir,
    setup_generation_log,
    log_event,
    claude_env
)
from lib.path_utils_v3 import get_next_iteration
from lib.feedback_history_v5 import get_feedback_history_for_claude


def execute(take_path: Path, story_name: str, inputs: dict, feedback=None, refined_from=None) -> dict:
    """
    Execute Stage 9: Script Adaptation.

    Args:
        take_path: Path to take directory
        story_name: Story name (slug)
        inputs: Dict mapping stage names to values, e.g.
                {"03-script": "006", "08-frames": {"01-scene": "003-refined-002", ...}}
        feedback: Optional user feedback text for this generation
        refined_from: Optional iteration number this is a refinement of

    Returns:
        Result dict with status, iteration, error (if any)
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 9: Script Adaptation")

    # Determine generation mode
    if refined_from and feedback:
        generation_mode = "refine"
        logger.info(f"Mode: REFINE (from {refined_from})")
    else:
        generation_mode = "new"
        logger.info(f"Mode: NEW (fresh adaptation)")

    # Read original script from 03-script/{inputs["03-script"]}
    script_iter = inputs.get("03-script")
    if not script_iter:
        return {
            "status": "error",
            "error": "No 03-script iteration specified in inputs"
        }

    script_iter_path = take_path / "03-script" / script_iter
    script_file = script_iter_path / "script.json"

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

    logger.info(f"Read original script: {script_file}")

    # Read selected frames from 08-frames
    frame_selections = inputs.get("08-frames")
    if not frame_selections or not isinstance(frame_selections, dict):
        return {
            "status": "error",
            "error": "No 08-frames selections in inputs. Select frames first."
        }

    frames_dir = take_path / "08-frames"
    frame_images = []  # [(scene_key, frame_name, frame_path), ...]
    scene_feedback = {}  # {scene_key: "feedback text"}

    for scene_key in sorted(frame_selections.keys()):
        # Skip title-card — it's handled by stage 11 (titles), not script adaptation
        if scene_key == "title-card":
            continue

        frame_name = frame_selections[scene_key]
        frame_path = frames_dir / scene_key / f"{frame_name}.jpg"

        if frame_path.exists():
            frame_images.append((scene_key, frame_name, frame_path))
            logger.info(f"Frame {scene_key}: {frame_name}.jpg")
        else:
            logger.warning(f"Frame not found: {frame_path}")

        # Collect refinement feedback
        base = frame_name.split("-refined-")[0] if "-refined-" in frame_name else frame_name
        fb_file = frames_dir / scene_key / f"feedback-{base}.txt"
        if fb_file.exists():
            scene_feedback[scene_key] = fb_file.read_text().strip()
            logger.info(f"Feedback for {scene_key}: {len(scene_feedback[scene_key])} chars")

    logger.info(f"Found {len(frame_images)} selected frames, {len(scene_feedback)} with feedback")

    # Determine iteration number
    stage_dir = take_path / "09-script-adapted"
    ensure_dir(stage_dir)

    iteration = get_next_iteration(stage_dir)
    iter_path = stage_dir / iteration
    ensure_dir(iter_path)

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

    # Save feedback.txt if feedback is 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}")

    # Check for feedback history
    feedback_history = get_feedback_history_for_claude(take_path, "09-script-adapted", refined_from)

    if feedback_history:
        logger.info("Found previous feedback")

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

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

    template = prompt_template_path.read_text()

    # Replace paths in template with actual relative paths
    story_path = take_path.parent
    rel_script_path = script_file.relative_to(story_path)

    prompt = template.replace("{{SCRIPT_PATH}}", str(rel_script_path))

    # Build selected frames section
    frames_section = "\n## Selected Frames\n\n"
    frames_section += "These are the ACTUAL frames selected and refined by the user.\n"
    frames_section += "Read each frame image with your vision capabilities.\n\n"

    for scene_key, frame_name, frame_path in frame_images:
        scene_num = int(scene_key.split('-')[0])
        rel_frame_path = frame_path.relative_to(story_path)
        is_refined = "-refined-" in frame_name
        frames_section += f"### Scene {scene_num} ({scene_key})\n"
        frames_section += f"**Frame**: `{rel_frame_path}`"
        if is_refined:
            frames_section += f" (refined from base)"
        frames_section += f"\n"

        # Add refinement feedback for this scene
        if scene_key in scene_feedback:
            frames_section += f"**Refinement feedback applied**:\n"
            for line in scene_feedback[scene_key].split('\n'):
                frames_section += f"> {line}\n"
        frames_section += "\n"

    prompt += frames_section

    # Add generation mode instructions
    if generation_mode == "new":
        prompt += f"\n\n## Generation Mode: NEW\n\n"
        prompt += f"This is a fresh script adaptation from the selected frames.\n\n"
        if feedback:
            prompt += f"### User Guidance\n\n{feedback}\n\n"
    elif generation_mode == "refine":
        prompt += f"\n\n## Generation Mode: REFINE\n\n"
        prompt += f"You are refining the adapted script from iteration **{refined_from}**.\n\n"
        if feedback_history:
            prompt += f"### Accumulated Feedback\n\n{feedback_history}\n\n"
        if feedback:
            prompt += f"### Current Feedback (HIGHEST PRIORITY)\n\n{feedback}\n\n"

    # Prepend output directory so Claude knows where to write files
    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}"

    # Save prompt
    prompt_file = iter_path / "prompt.txt"
    with open(prompt_file, 'w') as f:
        f.write(prompt)
    logger.info(f"Saved prompt to {prompt_file}")

    # Call Claude with vision capabilities
    logger.info("Calling Claude with vision...")
    print("\n" + "=" * 60)
    print(f"Invoking Claude for Stage 9: Script Adaptation")
    print(f"Session ID: {session_id}")
    print(f"Iteration: {iteration}")
    print(f"Frames: {len(frame_images)} selected frames")
    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=900,  # 15 minutes for vision + analysis
            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:
        error_msg = "Claude timed out after 15 minutes"
        logger.error(error_msg)
        return {
            "status": "error",
            "error": error_msg
        }

    except Exception as e:
        error_msg = f"Exception calling Claude: {str(e)}"
        logger.error(error_msg)
        return {
            "status": "error",
            "error": error_msg
        }

    # Validate output files
    adapted_script_file = iter_path / "script.json"
    if not adapted_script_file.exists():
        error_msg = f"Expected output file not created: script.json"
        logger.error(error_msg)
        return {
            "status": "error",
            "error": error_msg
        }

    logger.info(f"Created script.json ({adapted_script_file.stat().st_size} bytes)")

    # Check for user message
    user_message_file = iter_path / "user_message.txt"
    user_message = None
    if user_message_file.exists():
        user_message = user_message_file.read_text().strip()
        logger.info(f"Created user_message.txt ({user_message_file.stat().st_size} bytes)")
    else:
        logger.warning("user_message.txt not created by Claude")
        user_message = "Script adaptation complete. Script now matches selected frames."

    # Check for feedback_request.json
    feedback_request_file = iter_path / "feedback_request.json"
    if feedback_request_file.exists():
        logger.info(f"Created feedback_request.json ({feedback_request_file.stat().st_size} bytes)")
    else:
        logger.warning("feedback_request.json not created by Claude")

    # Create V5 metadata
    metadata = {
        "stage": "09-script-adapted",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": {
            "03-script": script_iter,
            "08-frames": frame_selections,
        },
        "refined_from": refined_from,
        "rating": None,
        "result": f"09-script-adapted/{iteration}/script.json",
        "model": "claude-sonnet-4.5",
        "session_id": session_id,
    }

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

    # Log to pipeline.log
    log_event(
        story_path,
        "stage_complete",
        {
            "action": "script_adapted",
            "iteration": iteration,
            "refined_from": refined_from,
            "result": f"09-script-adapted/{iteration}/script.json"
        },
        take=take_path.name,
        stage="09-script-adapted",
        iteration=iteration
    )
    logger.info("Logged to pipeline.log")

    print("\n" + "=" * 60)
    print("Stage 9: Adapted script generated successfully")
    print(f"Location: {iter_path / 'script.json'}")
    print("=" * 60)

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"09-script-adapted/{iteration}/script.json",
        "user_message": user_message
    }
