"""
长任务管理器 - 处理需要多步骤的复杂任务
"""
import asyncio
from typing import List, Dict, Optional, Callable, Any
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
import json
from pathlib import Path
import uuid


class TaskStatus(Enum):
    """任务状态"""
    PENDING = "pending"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class TaskStep:
    """任务步骤"""
    step_id: str
    description: str
    action: str
    params: Dict[str, Any]
    status: TaskStatus = TaskStatus.PENDING
    result: Optional[Any] = None
    error: Optional[str] = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    
    def to_dict(self) -> Dict:
        return {
            'step_id': self.step_id,
            'description': self.description,
            'action': self.action,
            'params': self.params,
            'status': self.status.value,
            'result': self.result,
            'error': self.error,
            'started_at': self.started_at.isoformat() if self.started_at else None,
            'completed_at': self.completed_at.isoformat() if self.completed_at else None
        }


@dataclass
class LongTask:
    """长任务"""
    task_id: str
    user_id: str
    title: str
    description: str
    steps: List[TaskStep] = field(default_factory=list)
    status: TaskStatus = TaskStatus.PENDING
    created_at: datetime = field(default_factory=datetime.now)
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    progress: float = 0.0
    metadata: Dict = field(default_factory=dict)
    
    def to_dict(self) -> Dict:
        return {
            'task_id': self.task_id,
            'user_id': self.user_id,
            'title': self.title,
            'description': self.description,
            'steps': [step.to_dict() for step in self.steps],
            'status': self.status.value,
            'created_at': self.created_at.isoformat(),
            'started_at': self.started_at.isoformat() if self.started_at else None,
            'completed_at': self.completed_at.isoformat() if self.completed_at else None,
            'progress': self.progress,
            'metadata': self.metadata
        }
    
    @classmethod
    def from_dict(cls, data: Dict) -> 'LongTask':
        steps = [
            TaskStep(
                step_id=s['step_id'],
                description=s['description'],
                action=s['action'],
                params=s['params'],
                status=TaskStatus(s['status']),
                result=s.get('result'),
                error=s.get('error'),
                started_at=datetime.fromisoformat(s['started_at']) if s.get('started_at') else None,
                completed_at=datetime.fromisoformat(s['completed_at']) if s.get('completed_at') else None
            )
            for s in data['steps']
        ]
        
        return cls(
            task_id=data['task_id'],
            user_id=data['user_id'],
            title=data['title'],
            description=data['description'],
            steps=steps,
            status=TaskStatus(data['status']),
            created_at=datetime.fromisoformat(data['created_at']),
            started_at=datetime.fromisoformat(data['started_at']) if data.get('started_at') else None,
            completed_at=datetime.fromisoformat(data['completed_at']) if data.get('completed_at') else None,
            progress=data.get('progress', 0.0),
            metadata=data.get('metadata', {})
        )


class LongTaskManager:
    """长任务管理器"""
    
    def __init__(self, storage_path: Path):
        self.storage_path = storage_path
        self.tasks: Dict[str, LongTask] = {}
        self.action_handlers: Dict[str, Callable] = {}
        self.running_tasks: Dict[str, asyncio.Task] = {}
        self._load()
    
    def _load(self):
        """加载任务"""
        if self.storage_path.exists():
            with open(self.storage_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                for task_data in data:
                    task = LongTask.from_dict(task_data)
                    self.tasks[task.task_id] = task
    
    def _save(self):
        """保存任务"""
        self.storage_path.parent.mkdir(parents=True, exist_ok=True)
        data = [task.to_dict() for task in self.tasks.values()]
        with open(self.storage_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    
    def register_action(self, action_name: str, handler: Callable):
        """注册动作处理器"""
        self.action_handlers[action_name] = handler
    
    def create_task(
        self,
        user_id: str,
        title: str,
        description: str,
        steps: List[Dict],
        metadata: Optional[Dict] = None
    ) -> str:
        """创建任务"""
        task_id = str(uuid.uuid4())
        
        task_steps = [
            TaskStep(
                step_id=str(uuid.uuid4()),
                description=step['description'],
                action=step['action'],
                params=step.get('params', {})
            )
            for step in steps
        ]
        
        task = LongTask(
            task_id=task_id,
            user_id=user_id,
            title=title,
            description=description,
            steps=task_steps,
            metadata=metadata or {}
        )
        
        self.tasks[task_id] = task
        self._save()
        
        return task_id
    
    async def execute_task(
        self,
        task_id: str,
        progress_callback: Optional[Callable[[str, float, str], None]] = None
    ) -> bool:
        """执行任务"""
        if task_id not in self.tasks:
            return False
        
        task = self.tasks[task_id]
        
        if task.status == TaskStatus.RUNNING:
            return False
        
        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now()
        self._save()
        
        try:
            total_steps = len(task.steps)
            
            for i, step in enumerate(task.steps):
                # 检查是否被取消
                if task.status == TaskStatus.CANCELLED:
                    break
                
                step.status = TaskStatus.RUNNING
                step.started_at = datetime.now()
                
                # 执行步骤
                try:
                    handler = self.action_handlers.get(step.action)
                    if not handler:
                        raise ValueError(f"未找到动作处理器: {step.action}")
                    
                    result = await handler(**step.params)
                    step.result = result
                    step.status = TaskStatus.COMPLETED
                    
                except Exception as e:
                    step.error = str(e)
                    step.status = TaskStatus.FAILED
                    task.status = TaskStatus.FAILED
                    self._save()
                    return False
                
                finally:
                    step.completed_at = datetime.now()
                
                # 更新进度
                task.progress = (i + 1) / total_steps
                self._save()
                
                if progress_callback:
                    await progress_callback(
                        task_id,
                        task.progress,
                        f"完成步骤 {i+1}/{total_steps}: {step.description}"
                    )
            
            task.status = TaskStatus.COMPLETED
            task.completed_at = datetime.now()
            task.progress = 1.0
            self._save()
            
            return True
            
        except Exception as e:
            task.status = TaskStatus.FAILED
            self._save()
            raise e
    
    async def start_task(
        self,
        task_id: str,
        progress_callback: Optional[Callable[[str, float, str], None]] = None
    ):
        """启动任务（异步）"""
        if task_id in self.running_tasks:
            return
        
        async def run():
            try:
                await self.execute_task(task_id, progress_callback)
            finally:
                if task_id in self.running_tasks:
                    del self.running_tasks[task_id]
        
        self.running_tasks[task_id] = asyncio.create_task(run())
    
    def pause_task(self, task_id: str) -> bool:
        """暂停任务"""
        if task_id not in self.tasks:
            return False
        
        task = self.tasks[task_id]
        if task.status == TaskStatus.RUNNING:
            task.status = TaskStatus.PAUSED
            self._save()
            return True
        
        return False
    
    def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        if task_id not in self.tasks:
            return False
        
        task = self.tasks[task_id]
        task.status = TaskStatus.CANCELLED
        self._save()
        
        # 取消正在运行的任务
        if task_id in self.running_tasks:
            self.running_tasks[task_id].cancel()
        
        return True
    
    def get_task(self, task_id: str) -> Optional[LongTask]:
        """获取任务"""
        return self.tasks.get(task_id)
    
    def get_user_tasks(self, user_id: str) -> List[LongTask]:
        """获取用户的所有任务"""
        return [
            task for task in self.tasks.values()
            if task.user_id == user_id
        ]
    
    def get_task_status(self, task_id: str) -> Optional[Dict]:
        """获取任务状态"""
        task = self.get_task(task_id)
        if not task:
            return None
        
        completed_steps = sum(1 for step in task.steps if step.status == TaskStatus.COMPLETED)
        total_steps = len(task.steps)
        
        return {
            'task_id': task.task_id,
            'title': task.title,
            'status': task.status.value,
            'progress': task.progress,
            'completed_steps': completed_steps,
            'total_steps': total_steps,
            'current_step': next(
                (step.description for step in task.steps if step.status == TaskStatus.RUNNING),
                None
            )
        }
