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

Reads story from 01-story/current/story.md, calls Claude to generate understanding,
creates 02-understanding/001/understanding.md, updates task queue.
"""

import argparse
import sys
from pathlib import Path
from datetime import datetime, timezone

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

from lib import (
    load_queue,
    save_queue,
    get_task_by_id,
    mark_completed,
    mark_waiting,
    append_feedback,
    get_feedback_history_for_claude,
    save_metadata,
    update_current_symlink,
    ensure_dir,
    log_event,
    setup_generation_log
)


def read_story(take_path: Path) -> tuple[str, Path]:
    """Read story from 01-story/current/story.md."""
    story_current = take_path / "01-story" / "current"
    if not story_current.is_symlink():
        raise ValueError("No current iteration set for 01-story")

    story_iter_path = story_current.resolve()
    story_file = story_iter_path / "story.md"

    if not story_file.exists():
        raise ValueError(f"Story file not found: {story_file}")

    story_text = story_file.read_text()
    return story_text, story_file


def call_claude_understanding(story_text: str, feedback_history: str) -> str:
    """
    Call Claude to generate understanding.md.

    For now, this is a placeholder that will use Claude Code's interactive mode.
    In automated mode, this would use the API directly.
    """
    # TODO: Implement Claude API call
    # For V4 development, we'll use Claude Code interactively
    raise NotImplementedError(
        "Stage 2 needs Claude interaction. "
        "This should be called from orchestrator which handles Claude prompts."
    )


def create_understanding(
    take_path: Path,
    story_name: str,
    understanding_text: str,
    parent_iteration: str = None
) -> Path:
    """
    Create 02-understanding iteration with understanding.md and metadata.

    Args:
        take_path: Path to take directory
        story_name: Story name for logging
        understanding_text: Generated understanding content
        parent_iteration: Previous iteration if this is a refinement

    Returns:
        Path to created iteration directory
    """
    logger = setup_generation_log(take_path)

    # Determine iteration number
    understanding_stage = take_path / "02-understanding"
    ensure_dir(understanding_stage)

    existing_iters = [
        d.name for d in understanding_stage.iterdir()
        if d.is_dir() and d.name.isdigit() and len(d.name) == 3
    ]

    if existing_iters:
        max_iter = max(int(i) for i in existing_iters)
        iteration = f"{max_iter + 1:03d}"
    else:
        iteration = "001"

    iter_path = understanding_stage / iteration
    ensure_dir(iter_path)

    logger.info(f"Creating 02-understanding/{iteration}")

    # Write understanding.md
    understanding_file = iter_path / "understanding.md"
    with open(understanding_file, 'w') as f:
        f.write(understanding_text)

    logger.info(f"Wrote {understanding_file}")

    # Create metadata
    story_file = take_path / "01-story" / "current" / "story.md"
    metadata = {
        "stage": "02-understanding",
        "take": "001",  # TODO: Get from take_path
        "iteration": iteration,
        "parent_iteration": parent_iteration,
        "source": str(story_file.relative_to(iter_path.parent.parent)),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "model": "claude-sonnet-4.5",  # TODO: Get from config
        "prompt_template": "prompts/stage2_understanding_v4.md"
    }

    save_metadata(iter_path / "metadata.json", metadata)
    logger.info(f"Saved metadata to {iter_path / 'metadata.json'}")

    # Update current symlink
    update_current_symlink(understanding_stage, iteration)
    logger.info(f"Updated current symlink to {iteration}")

    return iter_path


def run_stage2_generate(take_path: Path, story_name: str) -> None:
    """
    Execute Stage 2: Generate understanding.

    1. Read story from 01-story/current
    2. Check for previous feedback
    3. Call Claude to generate understanding
    4. Create 02-understanding/001/understanding.md
    5. Update task queue (mark task 3 completed, task 4 waiting)
    6. Log to pipeline.log
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 2: Understanding Generation")

    # Load task queue
    queue = load_queue(take_path)
    task = get_task_by_id(queue, 3)  # Task 3: understanding_generate

    if not task:
        logger.error("Task 3 (understanding_generate) not found in queue")
        return

    if task['status'] == 'completed':
        logger.info("Task 3 already completed, skipping")
        return

    # Read story
    logger.info("Reading story from 01-story/current")
    story_text, story_file = read_story(take_path)
    logger.info(f"Read {len(story_text)} characters from {story_file}")

    # Check for previous feedback
    feedback_history = get_feedback_history_for_claude(take_path, "02-understanding")

    if feedback_history and feedback_history != "No previous feedback for this stage.":
        logger.info("Found previous feedback for this stage")
    else:
        logger.info("No previous feedback")

    # For now, print the prompt and wait for manual Claude interaction
    print("\n" + "="*60)
    print("STAGE 2: Understanding Generation")
    print("="*60)
    print("\nStory:")
    print("-" * 60)
    print(story_text)
    print("-" * 60)

    if feedback_history and feedback_history != "No previous feedback for this stage.":
        print("\nFeedback History:")
        print("-" * 60)
        print(feedback_history)
        print("-" * 60)

    print("\nNext steps:")
    print("1. Read the story above")
    print("2. Generate understanding.md following the Stage 2 template")
    print("3. The understanding should include:")
    print("   - Retelling in your own words")
    print("   - Meaning/punchline")
    print("   - What makes it work")
    print("   - Visual potential")
    print("   - What won't work for video")
    print("\nOnce you have the understanding text, call:")
    print(f"  python stage2_understanding.py save <understanding_text>")
    print("\nOr provide understanding as a file:")
    print(f"  python stage2_understanding.py save --file understanding.md")


def run_stage2_save(take_path: Path, story_name: str,
                   understanding_text: str = None,
                   understanding_file: Path = None) -> None:
    """
    Save generated understanding and update task queue.

    Args:
        take_path: Path to take directory
        story_name: Story name for logging
        understanding_text: Understanding content (if provided directly)
        understanding_file: Path to file containing understanding (alternative)
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 2: Save Understanding")

    # Get understanding text
    if understanding_file:
        if not understanding_file.exists():
            logger.error(f"Understanding file not found: {understanding_file}")
            return
        understanding_text = understanding_file.read_text()
        logger.info(f"Read understanding from {understanding_file}")

    if not understanding_text:
        logger.error("No understanding text provided")
        return

    logger.info(f"Understanding length: {len(understanding_text)} characters")

    # Create understanding iteration
    iter_path = create_understanding(take_path, story_name, understanding_text)
    logger.info(f"Created {iter_path}")

    # Update task queue
    queue = load_queue(take_path)

    # Mark task 3 (understanding_generate) as completed
    task3 = get_task_by_id(queue, 3)
    if task3:
        task3['status'] = 'completed'
        task3['result'] = f"02-understanding/001/understanding.md"
        task3['iteration'] = "001"
        logger.info("Marked task 3 (understanding_generate) as completed")

    # Mark task 4 (understanding_await_feedback) as waiting
    task4 = get_task_by_id(queue, 4)
    if task4:
        task4['status'] = 'waiting'
        task4['message'] = "Understanding generated, awaiting user feedback"
        logger.info("Marked task 4 (understanding_await_feedback) as waiting")

    save_queue(take_path, queue)
    logger.info("Updated task queue")

    # Log to pipeline.log
    story_path = take_path.parent
    log_event(
        story_path,
        "stage_complete",
        {
            "action": "understanding_generated",
            "iteration": "001",
            "result": "02-understanding/001/understanding.md"
        },
        take="001",
        stage="02-understanding",
        iteration="001"
    )
    logger.info("Logged to pipeline.log")

    print("\n" + "="*60)
    print("Understanding saved successfully!")
    print("="*60)
    print(f"\nLocation: {iter_path}")
    print(f"\nNext steps:")
    print("1. Review the understanding at:")
    print(f"   {iter_path / 'understanding.md'}")
    print("2. Provide feedback:")
    print(f"   python stage2_understanding.py feedback --accept")
    print(f"   python stage2_understanding.py feedback --refine 'Make it darker'")
    print(f"   python stage2_understanding.py feedback --regenerate 'Focus more on X'")


def run_stage2_feedback(take_path: Path, story_name: str,
                       accept: bool = False,
                       refine: str = None,
                       regenerate: str = None) -> None:
    """
    Process user feedback for understanding.

    Args:
        take_path: Path to take directory
        story_name: Story name for logging
        accept: User accepts the understanding
        refine: User wants refinement (provide feedback text)
        regenerate: User wants complete regeneration (provide feedback text)
    """
    logger = setup_generation_log(take_path)
    logger.section("Stage 2: Process Feedback")

    if not any([accept, refine, regenerate]):
        logger.error("Must provide one of: --accept, --refine, --regenerate")
        return

    # Load queue
    queue = load_queue(take_path)
    task4 = get_task_by_id(queue, 4)

    if not task4 or task4['status'] != 'waiting':
        logger.error("Task 4 not in waiting state")
        return

    if accept:
        # User accepts - mark task 4 as completed
        task4['status'] = 'completed'
        task4['feedback'] = 'accepted'
        logger.info("User accepted understanding")

        # Log feedback
        append_feedback(
            take_path,
            stage="02-understanding",
            iteration="001",
            feedback_type="accept",
            text="User accepted understanding",
            action="proceed"
        )

        print("\n✓ Understanding accepted!")
        print("Ready to proceed to Stage 3 (Script Generation)")

    elif refine:
        # User wants refinement
        task4['status'] = 'pending'  # Back to pending so it regenerates
        task4['feedback'] = refine
        logger.info(f"User requested refinement: {refine}")

        # Log feedback
        append_feedback(
            take_path,
            stage="02-understanding",
            iteration="001",
            feedback_type="refine",
            text=refine,
            action="regenerate"
        )

        # Mark task 3 back to pending for regeneration
        task3 = get_task_by_id(queue, 3)
        if task3:
            task3['status'] = 'pending'

        print("\n✓ Feedback recorded!")
        print(f"Feedback: {refine}")
        print("\nRun stage2 generate again to create iteration 002 with feedback")

    elif regenerate:
        # User wants complete regeneration
        task4['status'] = 'pending'
        task4['feedback'] = regenerate
        logger.info(f"User requested regeneration: {regenerate}")

        # Log feedback
        append_feedback(
            take_path,
            stage="02-understanding",
            iteration="001",
            feedback_type="regenerate",
            text=regenerate,
            action="regenerate"
        )

        # Mark task 3 back to pending
        task3 = get_task_by_id(queue, 3)
        if task3:
            task3['status'] = 'pending'

        print("\n✓ Feedback recorded!")
        print(f"Feedback: {regenerate}")
        print("\nRun stage2 generate again to create iteration 002 with feedback")

    # Save updated queue
    save_queue(take_path, queue)
    logger.info("Updated task queue")

    # Log to pipeline.log
    story_path = take_path.parent
    log_event(
        story_path,
        "feedback",
        {
            "action": "user_feedback",
            "feedback_type": "accept" if accept else ("refine" if refine else "regenerate"),
            "feedback_text": refine or regenerate or "accepted"
        },
        take="001",
        stage="02-understanding",
        iteration="001"
    )


def main():
    parser = argparse.ArgumentParser(
        description="Stage 2: Story Understanding Generation"
    )
    parser.add_argument(
        "command",
        choices=["generate", "save", "feedback"],
        help="Command to execute"
    )
    parser.add_argument(
        "story_path",
        help="Path to story directory (e.g., /mnt/video/microdrama-v3/20260203-water-walking-dog)"
    )
    parser.add_argument(
        "understanding_text",
        nargs="?",
        help="Understanding text (for 'save' command)"
    )
    parser.add_argument(
        "--file",
        type=Path,
        help="Path to file containing understanding text (for 'save' command)"
    )
    parser.add_argument(
        "--accept",
        action="store_true",
        help="Accept the understanding (for 'feedback' command)"
    )
    parser.add_argument(
        "--refine",
        help="Request refinement with feedback (for 'feedback' command)"
    )
    parser.add_argument(
        "--regenerate",
        help="Request regeneration with feedback (for 'feedback' command)"
    )

    args = parser.parse_args()

    # Validate story path
    story_path = Path(args.story_path)
    if not story_path.exists():
        print(f"Error: Story path not found: {story_path}")
        sys.exit(1)

    # Get take path (use current symlink)
    take_link = story_path / "current"
    if not take_link.is_symlink():
        print(f"Error: No current take set for {story_path}")
        sys.exit(1)

    take_path = take_link.resolve()
    story_name = story_path.name.split('-', 1)[1] if '-' in story_path.name else story_path.name

    # Execute command
    if args.command == "generate":
        run_stage2_generate(take_path, story_name)

    elif args.command == "save":
        run_stage2_save(
            take_path,
            story_name,
            understanding_text=args.understanding_text,
            understanding_file=args.file
        )

    elif args.command == "feedback":
        run_stage2_feedback(
            take_path,
            story_name,
            accept=args.accept,
            refine=args.refine,
            regenerate=args.regenerate
        )


if __name__ == "__main__":
    main()
