#!/usr/bin/env python3
"""
Stage 2: Understanding Generation

Reads story from 01-story/001/story.md, calls Claude to analyze it,
creates 02-understanding/XXX/understanding.md.
"""

import subprocess
import sys
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: str = None, refined_from: str = None) -> dict:
    """
    Execute Stage 2: Understanding generation.

    Args:
        take_path: Path to take directory
        story_name: Story name (slug)
        inputs: Parent stage iterations (empty for stage 2)
        feedback: User feedback text (if any)
        refined_from: Iteration being refined (if refinement)
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 2: Understanding Generation")

    # Read story from 01-story (via pipeline inputs or fallback scan)
    story_iter = inputs.get("01-story")
    if story_iter:
        story_file = take_path / "01-story" / story_iter / "story.md"
        if not story_file.exists():
            return {"status": "error", "error": f"Story file not found: {story_file}"}
    else:
        # Fallback: scan 01-story/ for any story.md (legacy projects)
        story_dir = take_path / "01-story"
        story_file = None
        for d in sorted(story_dir.iterdir()):
            if d.is_dir() and d.name.isdigit():
                candidate = d / "story.md"
                if candidate.exists():
                    story_file = candidate
                    break

        if not story_file:
            return {"status": "error", "error": "Story file not found in 01-story/"}

    story_text = story_file.read_text()
    logger.info(f"Read story: {len(story_text)} characters")

    # Determine iteration number
    stage_dir = take_path / "02-understanding"
    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 if 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}")

    # Feedback history from refinement chain
    feedback_history = get_feedback_history_for_claude(
        take_path, "02-understanding", refined_from
    )

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

    if not prompt_template_path.exists():
        prompt = f"""# Stage 2: Story Understanding

You are analyzing a story for visual adaptation into a short-form video (45-60 seconds).

## Story

{story_text}

## Task

Write TWO files:

### 1. understanding.md
A detailed analysis including retelling, meaning, visual potential, challenges.

### 2. user_message.txt
A friendly message to the user explaining your analysis.
"""
    else:
        template = prompt_template_path.read_text()
        prompt = template.replace("{story}", story_text)
        prompt = prompt.replace("{story_name}", story_name)
        prompt = prompt.replace("{iteration}", iteration)

    # 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)")

    # Add feedback
    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
    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
    (iter_path / "prompt.txt").write_text(prompt)
    logger.info(f"Saved prompt")

    # Call Claude
    logger.info("Calling Claude...")
    print("\n" + "=" * 60)
    print(f"Stage 2: Understanding (iteration {iteration})")
    print(f"Session: {session_id}")
    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=600,
            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:
        return {"status": "error", "error": "Claude timed out after 10 minutes"}
    except Exception as e:
        return {"status": "error", "error": f"Exception calling Claude: {str(e)}"}

    # Validate output
    understanding_file = iter_path / "understanding.md"
    if not understanding_file.exists():
        return {"status": "error", "error": "Expected output: understanding.md not created"}

    logger.info(f"Created understanding.md ({understanding_file.stat().st_size} bytes)")

    # User message
    user_message_file = iter_path / "user_message.txt"
    user_message = (user_message_file.read_text().strip()
                    if user_message_file.exists()
                    else "Understanding analysis complete. Please review.")

    # Save metadata (V5 format)
    metadata = {
        "stage": "02-understanding",
        "iteration": iteration,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "based_on": inputs,
        "refined_from": refined_from,
        "rating": None,
        "result": f"02-understanding/{iteration}/understanding.md",
        "model": "claude-sonnet-4.5",
        "session_id": session_id,
        "video_generator": "kling-2.6",
    }

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

    # Log event
    log_event(
        story_path, "stage_complete",
        {"action": "understanding_generated", "iteration": iteration,
         "result": f"02-understanding/{iteration}/understanding.md"},
        take=take_path.name, stage="02-understanding", iteration=iteration
    )

    print(f"\nGenerated: 02-understanding/{iteration}/understanding.md")

    return {
        "status": "success",
        "iteration": iteration,
        "result": f"02-understanding/{iteration}/understanding.md",
        "user_message": user_message,
    }
