#!/usr/bin/env python3
"""
Task Queue Management for V4 Orchestrator

JSON-based task queue operations for managing workflow state.
Format: Single JSON file with array of tasks, not JSONL.

Task structure:
{
  "id": 1,
  "stage": "02-understanding",
  "action": "generate" | "await_feedback" | "extract_batch" | "generate_batch",
  "status": "pending" | "waiting" | "completed" | "failed",
  "deps": [0],  # List of task IDs that must complete first
  "iteration": "001",  # Current iteration for this stage
  "parent_iteration": null | "001",  # Which iteration this was based on
  "result": null | "path/to/output",  # Output path when completed
  "message": null | "User message",  # For await_feedback tasks
  "feedback": null | "User feedback text",  # Feedback text if provided
  "error": null | "Error message"  # Error if failed
}
"""

import json
from pathlib import Path
from typing import Optional, Dict, List, Any
from datetime import datetime


def load_queue(take_path: Path) -> List[Dict[str, Any]]:
    """
    Load task queue from take directory.

    Args:
        take_path: Path to take directory (e.g., /mnt/video/microdrama-v3/story/003/)

    Returns:
        List of task dictionaries, empty list if file doesn't exist
    """
    queue_file = take_path / "task_queue.json"

    if not queue_file.exists():
        return []

    with open(queue_file, 'r') as f:
        queue = json.load(f)

    return queue


def save_queue(take_path: Path, queue: List[Dict[str, Any]]) -> None:
    """
    Save task queue to take directory.

    Args:
        take_path: Path to take directory
        queue: List of task dictionaries
    """
    queue_file = take_path / "task_queue.json"

    with open(queue_file, 'w') as f:
        json.dump(queue, f, indent=2)


def get_next_task(queue: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
    """
    Get next executable task from queue.

    Returns task that is:
    - Status = 'pending'
    - All dependencies completed

    Returns None if no task available or if a task is waiting for feedback.

    Args:
        queue: List of task dictionaries

    Returns:
        Next task to execute, or None
    """
    # First check if any task is waiting for feedback
    for task in queue:
        if task['status'] == 'waiting':
            return None  # Block until feedback provided

    # Find next pending task with all deps completed
    for task in queue:
        if task['status'] != 'pending':
            continue

        # Check dependencies
        deps_complete = True
        for dep_id in task.get('deps', []):
            dep_task = next((t for t in queue if t['id'] == dep_id), None)
            if not dep_task:
                print(f"Warning: Task {task['id']} depends on non-existent task {dep_id}")
                deps_complete = False
                break
            if dep_task['status'] != 'completed':
                deps_complete = False
                break

        if deps_complete:
            return task

    return None


def get_waiting_task(queue: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
    """
    Get task that is currently waiting for feedback.

    Args:
        queue: List of task dictionaries

    Returns:
        Task with status='waiting', or None
    """
    for task in queue:
        if task['status'] == 'waiting':
            return task
    return None


def mark_completed(
    queue: List[Dict[str, Any]],
    task_id: int,
    result: Optional[str] = None,
    iteration: Optional[str] = None
) -> None:
    """
    Mark task as completed and optionally set result/iteration.

    Args:
        queue: List of task dictionaries
        task_id: ID of task to mark completed
        result: Optional output path
        iteration: Optional iteration number created
    """
    task = next((t for t in queue if t['id'] == task_id), None)
    if not task:
        raise ValueError(f"Task {task_id} not found in queue")

    task['status'] = 'completed'
    if result:
        task['result'] = result
    if iteration:
        task['iteration'] = iteration


def mark_waiting(
    queue: List[Dict[str, Any]],
    task_id: int,
    message: str
) -> None:
    """
    Mark task as waiting for user feedback.

    Args:
        queue: List of task dictionaries
        task_id: ID of task to mark waiting
        message: Message to display to user
    """
    task = next((t for t in queue if t['id'] == task_id), None)
    if not task:
        raise ValueError(f"Task {task_id} not found in queue")

    task['status'] = 'waiting'
    task['message'] = message


def mark_failed(
    queue: List[Dict[str, Any]],
    task_id: int,
    error: str
) -> None:
    """
    Mark task as failed with error message.

    Args:
        queue: List of task dictionaries
        task_id: ID of task to mark failed
        error: Error message
    """
    task = next((t for t in queue if t['id'] == task_id), None)
    if not task:
        raise ValueError(f"Task {task_id} not found in queue")

    task['status'] = 'failed'
    task['error'] = error


def provide_feedback(
    queue: List[Dict[str, Any]],
    task_id: int,
    feedback: str
) -> None:
    """
    Provide feedback to waiting task and mark it as pending.

    This allows the task to be re-executed with the feedback.

    Args:
        queue: List of task dictionaries
        task_id: ID of waiting task
        feedback: User feedback text
    """
    task = next((t for t in queue if t['id'] == task_id), None)
    if not task:
        raise ValueError(f"Task {task_id} not found in queue")

    if task['status'] != 'waiting':
        raise ValueError(f"Task {task_id} is not waiting for feedback (status: {task['status']})")

    task['feedback'] = feedback
    task['status'] = 'pending'  # Ready to execute with feedback


def add_task(
    queue: List[Dict[str, Any]],
    stage: str,
    action: str,
    deps: Optional[List[int]] = None,
    **kwargs
) -> int:
    """
    Add new task to queue.

    Args:
        queue: List of task dictionaries
        stage: Stage identifier (e.g., "02-understanding")
        action: Action type (generate/await_feedback/etc.)
        deps: List of task IDs this depends on
        **kwargs: Additional task fields

    Returns:
        Task ID of newly added task
    """
    # Get next task ID
    task_id = max([t['id'] for t in queue], default=0) + 1

    task = {
        'id': task_id,
        'stage': stage,
        'action': action,
        'status': 'pending',
        'deps': deps or [],
        'iteration': None,
        'parent_iteration': None,
        'result': None,
        'message': None,
        'feedback': None,
        'error': None
    }

    # Add any additional fields
    task.update(kwargs)

    queue.append(task)
    return task_id


def get_task_by_id(queue: List[Dict[str, Any]], task_id: int) -> Optional[Dict[str, Any]]:
    """
    Get task by ID.

    Args:
        queue: List of task dictionaries
        task_id: Task ID

    Returns:
        Task dictionary, or None if not found
    """
    return next((t for t in queue if t['id'] == task_id), None)


def get_tasks_by_stage(queue: List[Dict[str, Any]], stage: str) -> List[Dict[str, Any]]:
    """
    Get all tasks for a specific stage.

    Args:
        queue: List of task dictionaries
        stage: Stage identifier (e.g., "02-understanding")

    Returns:
        List of tasks for that stage
    """
    return [t for t in queue if t['stage'] == stage]


def count_tasks_by_status(queue: List[Dict[str, Any]]) -> Dict[str, int]:
    """
    Count tasks by status.

    Args:
        queue: List of task dictionaries

    Returns:
        Dict with counts: {'pending': 5, 'completed': 10, 'waiting': 1, 'failed': 0}
    """
    counts = {'pending': 0, 'completed': 0, 'waiting': 0, 'failed': 0}
    for task in queue:
        status = task['status']
        counts[status] = counts.get(status, 0) + 1
    return counts


def reset_task(queue: List[Dict[str, Any]], task_id: int) -> None:
    """
    Reset task to pending status, clearing result/error/feedback.

    Useful for re-running a failed task or regenerating with new feedback.

    Args:
        queue: List of task dictionaries
        task_id: ID of task to reset
    """
    task = next((t for t in queue if t['id'] == task_id), None)
    if not task:
        raise ValueError(f"Task {task_id} not found in queue")

    task['status'] = 'pending'
    task['result'] = None
    task['error'] = None
    task['feedback'] = None
