#!/usr/bin/env python3
"""
V4 Orchestrator - Non-interactive task execution with JSON feedback API.

Runs tasks from queue until hitting a feedback gate, then:
1. Writes feedback_request.json describing what feedback is needed
2. Exits cleanly
3. External system (Telegram bot) reads request, shows to user, writes feedback_response.json
4. Run again - reads response, continues execution

Usage:
    orchestrator_v4.py STORY_PATH

Example:
    orchestrator_v4.py /mnt/video/microdrama-v3/20260203-water-walking-dog
"""

import sys
import json
import time
from pathlib import Path

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

from lib import (
    load_queue,
    save_queue,
    get_next_task,
    get_task_by_id,
    append_feedback,
    setup_generation_log,
    update_current_symlink
)
from lib.cascade import (
    create_cascade_tasks,
    get_available_restart_stages,
    rewire_downstream_deps
)

# Import stage executors
from stages import (
    stage_02_understanding,
    stage_03_script,
    stage_04_assets,
    stage_05_references,
    stage_06_characteristic,
    stage_07_grids
)


# Stage executor registry
STAGE_EXECUTORS = {
    "02-understanding": {
        "generate": stage_02_understanding.execute
    },
    "03-script": {
        "generate": stage_03_script.execute
    },
    "04-assets": {
        "generate": stage_04_assets.execute
    },
    "05-references": {
        "check_skip": stage_05_references.check_skip,
        "search_or_skip": stage_05_references.search_or_skip
    },
    "06-characteristic": {
        "generate": stage_06_characteristic.execute
    },
    "07-grids": {
        "generate": stage_07_grids.execute
    },
    # More stages will be added as we implement them
}


def execute_task(take_path: Path, story_name: str, task: dict, logger) -> dict:
    """
    Execute a single task by routing to appropriate stage executor.

    Args:
        take_path: Path to take directory
        story_name: Story name (slug)
        task: Task dictionary from queue
        logger: Generation logger

    Returns:
        Result dict with status, iteration, error (if any)
    """
    stage = task["stage"]
    action = task["action"]

    logger.info(f"Executing task {task['id']}: {stage} / {action}")

    # Skip completed tasks
    if task["status"] == "completed":
        logger.info("Task already completed, skipping")
        return {"status": "skip"}

    # Handle init tasks (already completed)
    if stage == "00-init":
        return {"status": "skip"}

    # Handle story task (already completed at init)
    if stage == "01-story":
        return {"status": "skip"}

    # Route to stage executor
    if stage in STAGE_EXECUTORS and action in STAGE_EXECUTORS[stage]:
        executor = STAGE_EXECUTORS[stage][action]
        return executor(take_path, story_name, task)

    # Unknown stage/action
    return {
        "status": "error",
        "error": f"No executor for {stage}/{action}"
    }


def save_feedback_request(take_path: Path, task: dict, user_message: str, logger) -> None:
    """
    Save feedback request as JSON for external systems (Telegram bot, etc).

    Args:
        take_path: Path to take directory
        task: Task dictionary (await_feedback type)
        user_message: Claude-generated message to user
        logger: Generation logger
    """
    stage = task["stage"]
    message = task.get("message", f"Review {stage}")

    # Get generated file path and iteration directory
    stage_current = take_path / stage / "current"
    generated_file = None
    file_description = None
    iter_path = None

    if stage_current.is_symlink():
        current_iter = stage_current.resolve()
        iter_path = current_iter

        # Map stage to output files
        output_files = {
            "02-understanding": "understanding.md",
            "03-script": "script.json",
            "04-assets": "assets.json",
            "06-characteristic": "characteristic_shot.*",  # .png or .jpg
            "07-grids": "grid.jpg",
            "08-script-adapted": "script.json",
            "11-assembly": "final.mp4"
        }

        if stage in output_files:
            pattern = output_files[stage]
            if "*" in pattern:
                # Glob pattern - find matching file
                matches = list(current_iter.glob(pattern))
                if matches:
                    filepath = matches[0]
                    generated_file = str(filepath)
                    file_size = filepath.stat().st_size
                    file_description = f"{filepath.name} ({file_size} bytes)"
            else:
                filepath = current_iter / pattern
                if filepath.exists():
                    generated_file = str(filepath)
                    file_size = filepath.stat().st_size
                    file_description = f"{pattern} ({file_size} bytes)"

    # Build restart_from options (stages before current)
    restart_stages = get_available_restart_stages(stage)

    # Create feedback request JSON
    feedback_request = {
        "task_id": task["id"],
        "stage": stage,
        "iteration": task.get("iteration"),
        "message": message,
        "user_message": user_message,
        "generated_file": generated_file,
        "file_description": file_description,
        "options": [
            {
                "action": "accept",
                "label": "✓ Accept",
                "description": "Accept and proceed to next stage",
                "requires_text": False
            },
            {
                "action": "refine",
                "label": "✏️ Refine",
                "description": "Request refinement (regenerate with feedback)",
                "requires_text": True
            },
            {
                "action": "regenerate",
                "label": "🔄 Regenerate",
                "description": "Completely regenerate from scratch",
                "requires_text": True
            }
        ]
    }

    # Add restart_from option if there are earlier stages to go back to
    if restart_stages:
        feedback_request["options"].append({
            "action": "restart_from",
            "label": "⏪ Restart from...",
            "description": "Go back to an earlier stage and regenerate all stages from there",
            "requires_text": True,
            "requires_target_stage": True,
            "available_stages": restart_stages
        })

    # Save to feedback_request.json in iteration directory
    if not iter_path:
        logger.error(f"Cannot save feedback request - iteration path not found for {stage}")
        return

    feedback_request_file = iter_path / "feedback_request.json"
    with open(feedback_request_file, 'w') as f:
        json.dump(feedback_request, f, indent=2)

    logger.info(f"Saved feedback request to {feedback_request_file}")
    print("\n" + "=" * 60)
    print("FEEDBACK NEEDED")
    print("=" * 60)
    print(f"Stage: {stage}")
    print(f"Iteration: {iter_path.name}")
    print(f"Message: {message}")
    if generated_file:
        print(f"Generated: {file_description}")
        print(f"File: {generated_file}")
    print(f"\nFeedback request saved to: {feedback_request_file}")
    print("\nOptions:")
    for opt in feedback_request["options"]:
        print(f"  {opt['label']}: {opt['description']}")
    print("\nOrchestrator will resume when feedback_response.json is provided in same directory.")
    print("=" * 60)


def load_feedback_response(take_path: Path, waiting_task: dict, logger) -> dict:
    """
    Load feedback response from external system.

    Args:
        take_path: Path to take directory
        waiting_task: Task that's waiting for feedback
        logger: Generation logger

    Returns:
        Dict with {action, feedback_text} or None if no response available
    """
    # Get iteration directory for the waiting task
    stage = waiting_task["stage"]
    iteration = waiting_task.get("iteration")

    if not iteration:
        logger.warning(f"Waiting task {waiting_task['id']} has no iteration set")
        return None

    iter_path = take_path / stage / iteration
    if not iter_path.exists():
        logger.warning(f"Iteration path not found: {iter_path}")
        return None

    feedback_response_file = iter_path / "feedback_response.json"

    if not feedback_response_file.exists():
        return None

    with open(feedback_response_file, 'r') as f:
        response = json.load(f)

    logger.info(f"Loaded feedback response: {response['action']}")

    # Delete response file after reading (one-time use)
    feedback_response_file.unlink()

    return response


def process_feedback_response(take_path: Path, story_name: str, task: dict,
                              response: dict, queue: list, logger) -> None:
    """
    Process user feedback and update queue accordingly.

    Args:
        take_path: Path to take directory
        story_name: Story name
        task: Feedback task dictionary
        response: User response {action, feedback_text}
        queue: Task queue
        logger: Generation logger
    """
    stage = task["stage"]
    action = response["action"]
    feedback_text = response.get("feedback_text")

    logger.info(f"Processing feedback for task {task['id']}: {action}")

    if action == "accept":
        # Optional: accept a specific iteration (switch current symlink)
        accept_iteration = response.get("iteration")
        if accept_iteration:
            stage_dir = take_path / stage
            update_current_symlink(stage_dir, accept_iteration)
            logger.info(f"Switched {stage}/current to {accept_iteration}")

        # Mark task as completed
        task["status"] = "completed"
        task["feedback"] = "accepted"
        logger.info(f"Task {task['id']} accepted")

        # Log feedback
        append_feedback(
            take_path,
            stage=stage,
            iteration=task.get("iteration", "current"),
            feedback_type="accept",
            text="User accepted",
            action="proceed"
        )

    elif action in ["refine", "regenerate"]:
        if not feedback_text:
            logger.warning(f"No feedback text provided for {action}")
            feedback_text = f"User requested {action}"

        # Log feedback for the current iteration
        append_feedback(
            take_path,
            stage=stage,
            iteration=task.get("iteration", "current"),
            feedback_type=action,
            text=feedback_text,
            action="regenerate"
        )
        logger.info(f"User feedback ({action}): {feedback_text}")

        # Keep current tasks as completed (don't reset them)
        task["status"] = "completed"
        task["feedback"] = f"{action}: {feedback_text}"

        # Find corresponding generate task
        generate_task_id = task["id"] - 1
        generate_task = get_task_by_id(queue, generate_task_id)
        if generate_task:
            generate_task["status"] = "completed"
            generate_task["feedback"] = feedback_text

        # Create NEW tasks for the regeneration
        max_id = max(t["id"] for t in queue)

        # New generate task
        new_generate_task = {
            "id": max_id + 1,
            "stage": stage,
            "action": "generate",
            "status": "pending",
            "deps": [task["id"]],  # Depends on the feedback that triggered it
            "iteration": None,  # Will be set during execution
            "parent_iteration": task.get("iteration"),
            "result": None,
            "message": f"Regenerate based on feedback: {feedback_text[:50]}...",
            "feedback": None,
            "error": None
        }

        # New await_feedback task
        new_feedback_task = {
            "id": max_id + 2,
            "stage": stage,
            "action": "await_feedback",
            "status": "pending",
            "deps": [max_id + 1],
            "iteration": None,
            "parent_iteration": None,
            "result": None,
            "message": f"Review {stage} (iteration will be created)",
            "feedback": None,
            "error": None
        }

        # Insert new tasks after the current feedback task
        current_index = queue.index(task)
        queue.insert(current_index + 1, new_generate_task)
        queue.insert(current_index + 2, new_feedback_task)

        logger.info(f"Created new tasks {max_id + 1} and {max_id + 2} for regeneration")

        # Rewire downstream deps: pending tasks that depended on the old
        # feedback task should now depend on the new feedback task
        new_task_ids = {max_id + 1, max_id + 2}
        cascade_stage_names = {stage}
        rewire_downstream_deps(queue, cascade_stage_names, max_id + 2, new_task_ids, logger)

    elif action == "restart_from":
        # Cascade regeneration: go back to an earlier stage
        target_stage = response.get("target_stage")
        if not target_stage:
            logger.error("restart_from action missing target_stage")
            return
        if not feedback_text:
            feedback_text = f"Restart from {target_stage}"

        logger.info(f"restart_from: {target_stage} (current: {stage})")

        # Delegate to cascade module
        create_cascade_tasks(
            queue=queue,
            feedback_task=task,
            target_stage=target_stage,
            current_stage=stage,
            feedback_text=feedback_text,
            take_path=take_path,
            logger=logger
        )


def switch_iteration(story_path: Path, stage: str, iteration: str) -> int:
    """
    Switch a stage to a specific iteration and reset downstream tasks.

    1. Switches current symlink for the stage
    2. Finds the next downstream generate task and resets it to pending
    3. Clears any failed tasks in the downstream chain

    Args:
        story_path: Path to story directory
        stage: Stage name (e.g. "06-characteristic")
        iteration: Iteration to switch to (e.g. "006")

    Returns:
        Exit code: 0=success, 1=error
    """
    take_link = story_path / "current"
    if not take_link.is_symlink():
        print(f"Error: No current take set for {story_path}")
        return 1

    take_path = take_link.resolve()
    logger = setup_generation_log(take_path)
    logger.section(f"Switch Iteration: {stage} -> {iteration}")

    # Verify iteration exists
    stage_dir = take_path / stage
    iter_dir = stage_dir / iteration
    if not iter_dir.exists():
        print(f"Error: Iteration not found: {iter_dir}")
        return 1

    # Switch symlink
    update_current_symlink(stage_dir, iteration)
    print(f"Switched {stage}/current -> {iteration}")
    logger.info(f"Switched {stage}/current -> {iteration}")

    # Load queue and reset downstream
    queue = load_queue(take_path)

    # Find the await_feedback task for this stage that gates downstream work
    # and mark it completed (accepted with this iteration)
    for task in queue:
        if (task["stage"] == stage
                and task["action"] == "await_feedback"
                and task["status"] in ("waiting", "pending")):
            task["status"] = "completed"
            task["feedback"] = f"accepted (switched to iteration {iteration})"
            task["iteration"] = iteration
            logger.info(f"Marked task {task['id']} ({stage}/await_feedback) as completed")
            print(f"Marked task {task['id']} ({stage}/await_feedback) as completed")

    # Find downstream generate tasks that depend on this stage's feedback
    # and reset them to pending
    reset_count = 0
    for task in queue:
        if task["status"] == "failed":
            # Reset any failed task — it may be the downstream generate
            task["status"] = "pending"
            task["error"] = None
            task["parent_iteration"] = None
            task["message"] = None
            reset_count += 1
            logger.info(f"Reset failed task {task['id']} ({task['stage']}/{task['action']}) to pending")
            print(f"Reset task {task['id']} ({task['stage']}/{task['action']}) to pending")

    save_queue(take_path, queue)
    print(f"\nDone. Reset {reset_count} failed tasks. Run orchestrator to continue.")
    return 0


def run_orchestrator(story_path: Path) -> int:
    """
    Main orchestrator loop.

    Executes tasks until hitting a feedback gate or queue empty.

    Args:
        story_path: Path to story directory (e.g., /mnt/video/microdrama-v3/20260203-water-walking-dog)

    Returns:
        Exit code: 0=completed, 1=error, 2=awaiting feedback
    """
    # 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}")
        return 1

    take_path = take_link.resolve()

    # Extract story name (slug) from directory name
    story_dirname = story_path.name
    if '-' in story_dirname:
        story_name = story_dirname.split('-', 1)[1]
    else:
        story_name = story_dirname

    # Setup logger
    logger = setup_generation_log(take_path)
    logger.section("Orchestrator V4 - Task Execution")
    logger.info(f"Story: {story_name}")
    logger.info(f"Take: {take_path.name}")

    print("\n" + "=" * 60)
    print(f"Orchestrator V4: {story_name}")
    print(f"Take: {take_path.name}")
    print("=" * 60)

    # Check for feedback response first
    queue = load_queue(take_path)
    waiting_tasks = [t for t in queue if t["status"] == "waiting"]

    if waiting_tasks:
        waiting_task = waiting_tasks[0]
        feedback_response = load_feedback_response(take_path, waiting_task, logger)

        if feedback_response:
            # Process the feedback
            process_feedback_response(take_path, story_name, waiting_task,
                                     feedback_response, queue, logger)
            save_queue(take_path, queue)
            print(f"✓ Processed feedback: {feedback_response['action']}")

    # Main execution loop
    iteration_count = 0
    max_iterations = 100  # Safety limit

    while iteration_count < max_iterations:
        iteration_count += 1

        # Load current queue state
        queue = load_queue(take_path)
        if not queue:
            print("\nError: Empty queue")
            return 1

        # Get next executable task
        next_task = get_next_task(queue)

        if not next_task:
            # Check if any tasks are waiting for feedback
            waiting_tasks = [t for t in queue if t["status"] == "waiting"]

            if waiting_tasks:
                # Create feedback request and exit
                waiting_task = waiting_tasks[0]
                logger.info(f"Task {waiting_task['id']} waiting for feedback")
                user_message = waiting_task.get("user_message", "Generation complete. Please review.")
                save_feedback_request(take_path, waiting_task, user_message, logger)
                return 2  # Exit code 2 = awaiting feedback

            else:
                # No tasks waiting, no tasks executable - we're done
                logger.info("No more executable tasks")
                print("\n" + "=" * 60)
                print("✓ All tasks completed")
                print("=" * 60)
                return 0  # Exit code 0 = completed

        # Execute the task
        logger.info(f"Processing task {next_task['id']}: {next_task['stage']} / {next_task['action']}")
        print(f"\n→ Task {next_task['id']}: {next_task['stage']} / {next_task['action']}")

        result = execute_task(take_path, story_name, next_task, logger)

        if result["status"] == "skip":
            # Task was already completed or not applicable
            next_task["status"] = "completed"
            save_queue(take_path, queue)
            continue

        elif result["status"] == "success":
            # Mark task as completed
            next_task["status"] = "completed"
            next_task["result"] = result.get("result")
            next_task["iteration"] = result.get("iteration")
            logger.info(f"Task {next_task['id']} completed successfully")

            # If this was a generate task, mark corresponding await_feedback as waiting
            if next_task["action"] == "generate":
                # Find the corresponding await_feedback task (usually ID + 1)
                feedback_task_id = next_task["id"] + 1
                feedback_task = get_task_by_id(queue, feedback_task_id)

                if feedback_task and feedback_task["action"] == "await_feedback":
                    feedback_task["status"] = "waiting"
                    feedback_task["iteration"] = result.get("iteration")
                    feedback_task["user_message"] = result.get("user_message", "Generation complete. Please review.")
                    logger.info(f"Task {feedback_task_id} now waiting for user feedback")

            save_queue(take_path, queue)

        elif result["status"] == "error":
            # Mark task as failed
            next_task["status"] = "failed"
            next_task["error"] = result.get("error")
            logger.error(f"Task {next_task['id']} failed: {result.get('error')}")
            save_queue(take_path, queue)

            print(f"\n✗ Task {next_task['id']} FAILED")
            print(f"Error: {result.get('error')}")
            print("\nOrchestrator stopped due to task failure")
            return 1  # Exit code 1 = error

        # Small delay for readability
        time.sleep(0.5)

    if iteration_count >= max_iterations:
        logger.warning("Reached maximum iteration limit")
        print("\nWarning: Reached maximum iteration limit (safety stop)")
        return 1


def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  orchestrator_v4.py STORY_PATH                          # Run orchestrator")
        print("  orchestrator_v4.py STORY_PATH switch STAGE ITERATION   # Switch iteration")
        print()
        print("Examples:")
        print("  orchestrator_v4.py /mnt/video/microdrama-v3/20260203-water-walking-dog")
        print("  orchestrator_v4.py /mnt/video/microdrama-v3/20260203-water-walking-dog switch 06-characteristic 006")
        print()
        print("Exit codes:")
        print("  0 - All tasks completed successfully")
        print("  1 - Error occurred")
        print("  2 - Awaiting user feedback (feedback_request.json created)")
        sys.exit(1)

    story_path = Path(sys.argv[1])

    if not story_path.exists():
        print(f"Error: Story path not found: {story_path}")
        sys.exit(1)

    # Check for switch command
    if len(sys.argv) >= 4 and sys.argv[2] == "switch":
        if len(sys.argv) < 5:
            print("Usage: orchestrator_v4.py STORY_PATH switch STAGE ITERATION")
            sys.exit(1)
        stage = sys.argv[3]
        iteration = sys.argv[4]
        exit_code = switch_iteration(story_path, stage, iteration)
        sys.exit(exit_code)

    exit_code = run_orchestrator(story_path)
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
