#!/usr/bin/env python3
"""
Stage 3: Script Generation (V5 orchestrator)

Reads understanding from 02-understanding/{inputs["02-understanding"]}/understanding.md,
calls Claude to create scene-by-scene script with dialogue, audio, and caption strategy.

Creates 03-script/XXX/script.json.

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

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
import reference_loader


def execute(take_path: Path, story_name: str, inputs: dict, feedback=None, refined_from=None, **kwargs) -> dict:
    """
    Execute Stage 3: Script generation.

    Args:
        take_path: Path to take directory
        story_name: Story name (slug)
        inputs: Dict mapping stage names to iteration numbers, e.g. {"02-understanding": "001"}
        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 3: Script Generation")

    # Read understanding from 02-understanding/{inputs["02-understanding"]}
    understanding_iter = inputs.get("02-understanding")
    if not understanding_iter:
        return {
            "status": "error",
            "error": "No 02-understanding iteration specified in inputs"
        }

    understanding_iter_path = take_path / "02-understanding" / understanding_iter
    understanding_file = understanding_iter_path / "understanding.md"

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

    understanding_text = understanding_file.read_text()
    logger.info(f"Read understanding: {len(understanding_text)} characters")

    # Load video generator model and video mode from understanding metadata
    understanding_metadata_file = take_path / "02-understanding" / understanding_iter / "metadata.json"
    video_generator = "kling-2.6"  # Default
    video_mode = "standard"  # Default
    if understanding_metadata_file.exists():
        with open(understanding_metadata_file, 'r') as f:
            understanding_metadata = json.load(f)
            video_generator = understanding_metadata.get("video_generator", "kling-2.6")
            video_mode = understanding_metadata.get("video_mode", "standard")

    logger.info(f"Video generator: {video_generator}")
    logger.info(f"Video mode: {video_mode}")

    # Determine clip duration constraints based on model
    if video_generator == "kling-2.6":
        duration_constraint = "5 or 10 seconds only"
        duration_note = "Use 5s for most scenes, 10s for scenes with longer dialogue or complex action"
    elif video_generator == "veo3":
        duration_constraint = "8 seconds only"
        duration_note = "All scenes must be exactly 8 seconds"
    else:
        duration_constraint = "5 or 10 seconds"  # Fallback to Kling
        duration_note = "Use 5s for most scenes, 10s for longer scenes"

    # Determine iteration number
    script_stage = take_path / "03-script"
    ensure_dir(script_stage)

    iteration = get_next_iteration(script_stage)
    iter_path = script_stage / 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, "03-script", refined_from)

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

    # Build prompt
    prompt_template_path = Path(__file__).parent.parent.parent / "prompts" / "stage_script.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()

    # Prepend video generator constraints to template
    prompt = f"""# Video Generator Model Constraints

**Model**: {video_generator}
**Clip durations**: {duration_constraint}
**Important**: {duration_note}

---

{template}"""

    # Add interpolation mode constraints if configured
    if video_mode == "interpolation":
        prompt += """

## INTERPOLATION VIDEO MODE

This project uses frame-pair interpolation to generate videos. Instead of independent video clips, we pair frames: scenes 1+2 form clip 1, scenes 3+4 form clip 2.

**Critical constraints for scene writing:**

1. **Scenes are PAIRED**: (1,2) and (3,4). Each pair will become ONE interpolated video clip.
2. **Each scene describes a STATIC STATE** (a freeze-frame moment), NOT an action or motion.
3. **Odd scenes (1,3) = START states**. Even scenes (2,4) = END states.
4. **Maximize visual CONTRAST within pairs** — dramatically different state between start and end.
5. **PREFER DISCONTINUITY between paired scenes.** End scenes (2,4) should use DIFFERENT camera angles, DIFFERENT locations or sections of the same space, and DIFFERENT framing from their start scenes (1,3). The AI will create a jarring spatial jump between frames — this is intentional and desirable. Examples:
   - Wide behind-shot of man in aisle → overhead bird's-eye of the same man in a different store section
   - Eye-level medium shot of one character → wide establishing shot revealing a crowd of them
   - Interior close-up → exterior wide shot of the same building
   Only use same-angle same-location pairs when the story specifically requires a smooth transformation (e.g., before/after of the exact same object).
6. **Dialogue goes in END scenes (even)** — any dialogue that happens during the transition belongs to the end scene. Start scenes (odd) should have `dialogue: null`.
7. **Add an `interpolation_prompt` field** to each scene in the JSON — a brief transition hint describing the spatial/perspective jump (e.g., "camera teleports to overhead, completely different section", "wide reveal — pull back to show the full scale").
"""
        logger.info("Added interpolation mode constraints")

    # Add template/references context if project has them
    project_path = take_path.parent
    template_context = reference_loader.get_full_prompt_context(project_path)
    if template_context:
        prompt += "\n\n" + template_context
        logger.info(f"Added template context ({len(template_context)} chars)")

    # Inject creative_notes from story_pipeline.json
    story_pipeline_file = project_path / "story_pipeline.json"
    if story_pipeline_file.exists():
        try:
            sp_data = json.loads(story_pipeline_file.read_text())
            creative_notes = sp_data.get("creative_notes")
            if creative_notes:
                prompt += (f"\n\n## CREATIVE DIRECTION\n\n"
                           f"Channel-level creative requirements:\n\n{creative_notes}\n")
                logger.info(f"Added creative_notes ({len(creative_notes)} chars)")
        except (json.JSONDecodeError, OSError):
            pass

    # Add feedback context
    if feedback_history:
        prompt += f"\n\n## Previous Feedback\n\n{feedback_history}\n"
    elif feedback:
        prompt += f"\n\n## User Feedback\n\n{feedback}\n"

    # Prepend output directory so Claude knows where to write files
    # cwd is story_path for session continuity, so Claude needs explicit path
    story_path = take_path.parent
    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
    logger.info("Calling Claude...")
    print("\n" + "=" * 60)
    print(f"Invoking Claude for Stage 3: Script")
    print(f"Session ID: {session_id}")
    print(f"Iteration: {iteration}")
    print("=" * 60)

    try:
        cmd = [
                "claude",
                "--resume", session_id,
                "--allowedTools", "Read,Write,Glob,Grep",
        ]
        model = _get_project_claude_model(story_path)
        if model:
            cmd.extend(["--model", model])
        cmd.extend(["-p", prompt])
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            cwd=str(story_path),
            timeout=600,  # 10 minutes
            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 10 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
    script_file = iter_path / "script.json"
    if not 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 ({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 generation complete. Please review the script.json file."

    # Create V5 metadata
    metadata = {
        "stage": "03-script",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"03-script/{iteration}/script.json",
        "model": "claude-sonnet-4.5",
        "session_id": session_id,
        "video_generator": video_generator,
        "video_mode": video_mode,
        "duration_constraint": duration_constraint,
    }

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

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

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

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


def _get_project_claude_model(story_path: Path) -> str:
    """Read claude_model from story_pipeline.json, or None for default."""
    pipeline_file = story_path / "story_pipeline.json"
    if pipeline_file.exists():
        try:
            data = json.loads(pipeline_file.read_text())
            return data.get("claude_model")
        except (json.JSONDecodeError, OSError):
            pass
    return None
