#!/usr/bin/env python3
"""
Run a stage of the microdrama workflow with Claude.

Usage:
    run_stage_v3.py STORY_NAME --stage STAGE [--take TAKE] [--feedback "text"]

Examples:
    # Run understanding stage
    run_stage_v3.py "fire-truck-brakes" --stage understanding

    # Run with feedback (creates new iteration)
    run_stage_v3.py "fire-truck-brakes" --stage understanding --feedback "More comedic tone"

    # Run specific take
    run_stage_v3.py "fire-truck-brakes" --stage script --take 001
"""
import argparse
import subprocess
import sys
import json
from pathlib import Path
from datetime import datetime

# Add lib to path
sys.path.insert(0, str(Path(__file__).parent))

from lib import (
    get_story_path,
    get_take_path,
    get_stage_path,
    get_stage_dirs,
    get_next_iteration,
    ensure_dir,
    update_current_symlink,
    log_stage_start,
    log_stage_complete,
    log_stage_error,
    log_feedback,
    save_metadata,
    create_metadata,
    save_prompt
)


# Stage configuration - will be expanded as we build each stage
STAGE_CONFIG = {
    "understanding": {
        "dir": "02-understanding",
        "prompt_template": "prompts/stage_understanding.txt",
        "output_file": "understanding.md",
        "description": "Analyze story for visual adaptation"
    },
    "script": {
        "dir": "03-script",
        "prompt_template": "prompts/stage_script.txt",
        "output_file": "script.json",
        "description": "Create scene-by-scene script with dialogue"
    },
    "assets": {
        "dir": "04-assets",
        "prompt_template": "prompts/stage_assets.txt",
        "output_file": "assets.json",
        "description": "Identify characters, locations, objects and reference needs"
    },
    "characteristic": {
        "dir": "06-characteristic",
        "prompt_template": "prompts/stage_characteristic.txt",
        "output_file": "nanobanana_prompt.txt",
        "description": "Create NanoBanana prompt for characteristic reference image"
    },
    "grid": {
        "dir": "07-grids",
        "prompt_template": "prompts/stage_grid.txt",
        "output_file": "nanobanana_prompt.txt",
        "description": "Create NanoBanana prompt for grid of scene reference frames"
    },
    "script-adapted": {
        "dir": "08-script-adapted",
        "prompt_template": "prompts/stage_script_adapted.txt",
        "output_file": "script.json",
        "description": "Adapt script to match generated grid frames"
    },
    # More stages will be added as we build them
}


def load_prompt_template(template_path: str, context: dict) -> str:
    """
    Load and fill prompt template.

    Args:
        template_path: Path to template file
        context: Dictionary of variables to fill in template

    Returns:
        Filled prompt text
    """
    template_file = Path(__file__).parent.parent / template_path

    if not template_file.exists():
        print(f"Warning: Template not found at {template_file}")
        print(f"Using basic prompt.")
        return f"Complete the {context.get('stage', 'task')}."

    template = template_file.read_text()

    # Simple template filling - can be enhanced later
    for key, value in context.items():
        template = template.replace(f"{{{key}}}", str(value))

    return template


def update_claude_md(story_name: str, take: str, stage: str, session_id: str) -> None:
    """
    Update CLAUDE.md with current state.

    Args:
        story_name: Story name
        take: Take number
        stage: Current stage
        session_id: Session ID
    """
    story_path = get_story_path(story_name)
    claude_md_path = story_path / "CLAUDE.md"

    if not claude_md_path.exists():
        print(f"Warning: CLAUDE.md not found, skipping update")
        return

    # Read current content
    content = claude_md_path.read_text()

    # Update current state section
    # This is a simple replacement - could be more sophisticated
    content = content.replace(
        f"- **Take**: {take}",
        f"- **Take**: {take}"
    )
    content = content.replace(
        "- **Stage**: (not started)",
        f"- **Stage**: {stage}"
    )
    content = content.replace(
        f"- **Session ID**: {story_name.replace(' ', '_')}_{take}_stage",
        f"- **Session ID**: {session_id}"
    )

    claude_md_path.write_text(content)


def run_stage(
    story_name: str,
    stage: str,
    take: str = "current",
    feedback: str = None
) -> bool:
    """
    Run a stage with Claude orchestration.

    Args:
        story_name: Story name
        stage: Stage name (e.g., "understanding", "script")
        take: Take number or "current"
        feedback: User feedback for iteration

    Returns:
        True if successful, False otherwise
    """
    # Validate stage
    if stage not in STAGE_CONFIG:
        print(f"Error: Unknown stage '{stage}'")
        print(f"Available stages: {', '.join(STAGE_CONFIG.keys())}")
        return False

    config = STAGE_CONFIG[stage]

    # Get paths
    try:
        story_path = get_story_path(story_name)
        take_path = get_take_path(story_name, take)
        take_num = take_path.name
    except ValueError as e:
        print(f"Error: {e}")
        return False

    stage_dir = config["dir"]
    stage_path = take_path / stage_dir

    # Ensure stage directory exists
    ensure_dir(stage_path)

    # Get next iteration
    iteration = get_next_iteration(stage_path)
    iter_path = stage_path / iteration
    ensure_dir(iter_path)

    print(f"Running stage: {stage} (iteration {iteration})")
    print(f"Description: {config['description']}")

    # Load session ID (one per story, reused for all stages/takes)
    session_file = story_path / "session.txt"
    if not session_file.exists():
        print(f"Error: Session file not found: {session_file}")
        print(f"Run init_take_v3.py first to initialize story")
        return False

    session_id = session_file.read_text().strip()
    print(f"Session ID: {session_id}")

    # Load prompt template
    context = {
        "story_name": story_name,
        "take": take_num,
        "stage": stage,
        "iteration": iteration,
        "session_id": session_id
    }

    # Add feedback if provided
    if feedback:
        context["feedback"] = feedback

    prompt = load_prompt_template(config["prompt_template"], context)

    # Add feedback to prompt if provided
    if feedback:
        prompt += f"\n\n## USER FEEDBACK\n\n{feedback}\n"

    # Save prompt
    save_prompt(iter_path / "prompt.txt", prompt)
    print(f"✓ Wrote prompt to {iter_path / 'prompt.txt'}")

    # Save feedback if provided
    if feedback:
        feedback_file = iter_path / "feedback.md"
        feedback_file.write_text(f"# User Feedback\n\n{feedback}\n")
        print(f"✓ Saved feedback to {feedback_file}")

    # Update CLAUDE.md
    update_claude_md(story_name, take_num, stage, session_id)

    # Log stage start
    log_stage_start(story_name, take_num, stage, iteration, session_id)

    # Invoke Claude
    print(f"\nInvoking Claude with session ID: {session_id}")
    print(f"Working directory: {iter_path}")
    print(f"\nPrompt preview:")
    print("-" * 60)
    print(prompt[:500] + ("..." if len(prompt) > 500 else ""))
    print("-" * 60)

    try:
        # Run Claude with -p for scripting, but allow tools for file operations
        result = subprocess.run(
            [
                "claude",
                "--session-id", session_id,
                "--allowedTools", "Read,Write,Glob,Grep",  # Tools Claude needs
                "-p", prompt
            ],
            capture_output=True,
            text=True,
            cwd=str(iter_path),  # Change working directory for subprocess
            timeout=600  # 10 minute timeout
        )

        # Check Claude output
        if result.returncode != 0:
            print(f"\nError: Claude returned code {result.returncode}")
            print(f"Stderr: {result.stderr}")
            log_stage_error(story_name, take_num, stage, iteration, session_id, result.stderr)
            return False

        print(f"\n✓ Claude completed")

    except subprocess.TimeoutExpired:
        error_msg = "Claude timed out after 10 minutes"
        print(f"\nError: {error_msg}")
        log_stage_error(story_name, take_num, stage, iteration, session_id, error_msg)
        return False

    except Exception as e:
        error_msg = str(e)
        print(f"\nError invoking Claude: {error_msg}")
        log_stage_error(story_name, take_num, stage, iteration, session_id, error_msg)
        return False

    # Validate output file exists
    output_file = iter_path / config["output_file"]
    if not output_file.exists():
        error_msg = f"Expected output file not created: {output_file}"
        print(f"\nError: {error_msg}")
        log_stage_error(story_name, take_num, stage, iteration, session_id, error_msg)
        return False

    print(f"✓ Output file created: {output_file}")

    # Create metadata
    metadata = create_metadata(
        stage=stage,
        take=take_num,
        iteration=iteration,
        session_id=session_id,
        prompt_template=config["prompt_template"],
        model="claude-sonnet-4.5",
        output_file=config["output_file"]
    )

    if feedback:
        metadata["feedback"] = {
            "from_user": feedback,
            "timestamp": datetime.now().isoformat()
        }

    # Save metadata
    save_metadata(iter_path / "metadata.json", metadata)
    print(f"✓ Saved metadata")

    # Update current symlink
    update_current_symlink(stage_path, iteration)
    print(f"✓ Updated current → {iteration}")

    # Log feedback if provided
    if feedback:
        log_feedback(story_name, take_num, stage, iteration, feedback)

    # Log completion
    log_stage_complete(story_name, take_num, stage, iteration, session_id, str(output_file))

    # Auto-generate image for stages that produce NanoBanana prompts
    if stage in ["grid", "characteristic"]:
        print(f"\n{'='*60}")
        print(f"Auto-generating image from prompt...")
        print(f"{'='*60}\n")

        # Determine which image generation script to use
        if stage == "grid":
            gen_script = Path(__file__).parent / "generate_grid_image_v3.py"
        else:  # characteristic
            gen_script = Path(__file__).parent / "generate_characteristic_image_v3.py"

        # Run image generation
        try:
            result = subprocess.run(
                [
                    sys.executable,
                    str(gen_script),
                    story_name,
                    "--take", take_num,
                    "--iteration", iteration
                ],
                capture_output=True,
                text=True,
                timeout=300
            )

            if result.returncode != 0:
                print(f"Warning: Image generation failed:")
                print(result.stderr)
                print(f"\nYou can manually run:")
                print(f"  python {gen_script.name} '{story_name}' --iteration {iteration}")
            else:
                print(result.stdout)

        except subprocess.TimeoutExpired:
            print(f"Warning: Image generation timed out")
            print(f"\nYou can manually run:")
            print(f"  python {gen_script.name} '{story_name}' --iteration {iteration}")
        except Exception as e:
            print(f"Warning: Image generation error: {e}")
            print(f"\nYou can manually run:")
            print(f"  python {gen_script.name} '{story_name}' --iteration {iteration}")

    print(f"\n{'='*60}")
    print(f"Stage '{stage}' completed successfully!")
    print(f"{'='*60}")
    print(f"Output: {output_file}")
    print(f"Iteration: {iteration}")
    print(f"\nNext steps:")
    print(f"1. Review the output file")
    print(f"2. If good, run next stage")
    print(f"3. If needs changes, run again with --feedback")

    return True


def main():
    parser = argparse.ArgumentParser(
        description="Run a stage of the microdrama workflow",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s "fire-truck-brakes" --stage understanding
  %(prog)s "fire-truck-brakes" --stage script --take 001
  %(prog)s "fire-truck-brakes" --stage understanding --feedback "More comedic"
        """
    )

    parser.add_argument(
        "story_name",
        help="Story name"
    )
    parser.add_argument(
        "--stage",
        required=True,
        choices=list(STAGE_CONFIG.keys()),
        help="Stage to run"
    )
    parser.add_argument(
        "--take",
        default="current",
        help="Take number (default: current)"
    )
    parser.add_argument(
        "--feedback",
        help="User feedback (creates new iteration)"
    )

    args = parser.parse_args()

    success = run_stage(
        args.story_name,
        args.stage,
        args.take,
        args.feedback
    )

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
