"""
任务队列管理器
管理并行任务（生图、动态表格维护等）的状态和生命周期
"""
import asyncio
from typing import Dict, List, Optional
from enum import Enum
from datetime import datetime


class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "pending"      # 等待中
    RUNNING = "running"      # 进行中
    COMPLETED = "completed"  # 已完成
    FAILED = "failed"        # 失败
    CANCELLED = "cancelled"  # 已取消


class TaskType(Enum):
    """任务类型枚举"""
    IMAGE_WORKFLOW = "image_workflow"
    DYNAMIC_TABLE = "dynamic_table"


class TaskItem:
    """任务项"""
    
    def __init__(self, task_id: str, task_type: TaskType, chat_id: str):
        self.task_id = task_id
        self.task_type = task_type
        self.chat_id = chat_id
        self.status = TaskStatus.PENDING
        self.created_at = datetime.now()
        self.started_at = None
        self.completed_at = None
        self.error = None
        self.metadata = {}  # 用于存储提示词、修改内容等
        
    def to_dict(self):
        """转换为字典格式（前端友好）"""
        return {
            "taskId": self.task_id,
            "taskType": self.task_type.value,
            "chatId": self.chat_id,
            "status": self.status.value,
            "createdAt": self.created_at.isoformat(),
            "startedAt": self.started_at.isoformat() if self.started_at else None,
            "completedAt": self.completed_at.isoformat() if self.completed_at else None,
            "error": self.error,
            "metadata": self.metadata
        }


class TaskQueueManager:
    """
    全局任务队列管理器
    
    功能：
    - 管理所有并行任务的生命周期
    - 支持按聊天ID查询任务
    - 支持取消任务
    - 自动清理已完成的任务
    """
    
    def __init__(self):
        self.tasks: Dict[str, TaskItem] = {}
        self.chat_tasks: Dict[str, List[str]] = {}  # chat_id -> [task_ids]
        self._lock = asyncio.Lock()
        
    async def add_task(self, task_id: str, task_type: TaskType, chat_id: str) -> TaskItem:
        """添加任务到队列"""
        async with self._lock:
            task = TaskItem(task_id, task_type, chat_id)
            self.tasks[task_id] = task
            
            if chat_id not in self.chat_tasks:
                self.chat_tasks[chat_id] = []
            self.chat_tasks[chat_id].append(task_id)
            
            return task
    
    async def start_task(self, task_id: str):
        """标记任务开始执行"""
        async with self._lock:
            if task_id in self.tasks:
                self.tasks[task_id].status = TaskStatus.RUNNING
                self.tasks[task_id].started_at = datetime.now()
    
    async def complete_task(self, task_id: str, metadata: dict = None):
        """标记任务完成"""
        async with self._lock:
            if task_id in self.tasks:
                self.tasks[task_id].status = TaskStatus.COMPLETED
                self.tasks[task_id].completed_at = datetime.now()
                if metadata:
                    self.tasks[task_id].metadata.update(metadata)
    
    async def fail_task(self, task_id: str, error: str):
        """标记任务失败"""
        async with self._lock:
            if task_id in self.tasks:
                self.tasks[task_id].status = TaskStatus.FAILED
                self.tasks[task_id].completed_at = datetime.now()
                self.tasks[task_id].error = error
    
    async def cancel_task(self, task_id: str) -> bool:
        """
        取消任务
        
        Returns:
            bool: 是否成功取消
        """
        async with self._lock:
            if task_id in self.tasks:
                task = self.tasks[task_id]
                if task.status in [TaskStatus.PENDING, TaskStatus.RUNNING]:
                    task.status = TaskStatus.CANCELLED
                    task.completed_at = datetime.now()
                    return True
        return False
    
    async def get_chat_tasks(self, chat_id: str, include_completed: bool = False) -> List[dict]:
        """
        获取某个聊天的所有任务
        
        Args:
            chat_id: 聊天ID
            include_completed: 是否包含已完成的任务
            
        Returns:
            List[dict]: 任务列表
        """
        async with self._lock:
            task_ids = self.chat_tasks.get(chat_id, [])
            tasks = []
            for task_id in task_ids:
                if task_id in self.tasks:
                    task = self.tasks[task_id]
                    # 根据参数决定是否包含已完成的任务
                    if include_completed or task.status in [TaskStatus.PENDING, TaskStatus.RUNNING]:
                        tasks.append(task.to_dict())
            return tasks
    
    async def cleanup_completed_tasks(self, chat_id: str):
        """清理已完成的任务"""
        async with self._lock:
            if chat_id in self.chat_tasks:
                task_ids = self.chat_tasks[chat_id]
                completed_ids = [
                    tid for tid in task_ids 
                    if tid in self.tasks and 
                    self.tasks[tid].status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]
                ]
                for tid in completed_ids:
                    del self.tasks[tid]
                    self.chat_tasks[chat_id].remove(tid)


# 全局实例
task_queue_manager = TaskQueueManager()