"""
异步任务管理器
负责任务的提交、调度、状态管理
"""
import asyncio
import uuid
import time
import os
from typing import Dict, List, Optional, Callable, Awaitable
from .task_store import TaskStore


# 任务状态
STATUS_PENDING = "pending"
STATUS_RUNNING = "running"
STATUS_COMPLETED = "completed"
STATUS_FAILED = "failed"
STATUS_CANCELLED = "cancelled"


class TaskManager:
    def __init__(self, store_dir: str = "./data/tasks", max_concurrent: int = 3):
        self.store = TaskStore(store_dir)
        self.max_concurrent = max_concurrent
        self._tasks: Dict[str, dict] = {}
        self._running = 0
        self._semaphore = asyncio.Semaphore(max_concurrent)
        self._worker_task: Optional[asyncio.Task] = None
        self._queue: asyncio.Queue = asyncio.Queue()
        self._executor: Optional[Callable[[dict], Awaitable[dict]]] = None

    async def start(self, executor: Callable[[dict], Awaitable[dict]]):
        """启动任务管理器，executor 是实际执行任务的异步函数"""
        self._executor = executor
        self._worker_task = asyncio.create_task(self._worker_loop())
        print(f"[TaskManager] 已启动，最大并发: {self.max_concurrent}")

    async def stop(self):
        """停止任务管理器"""
        if self._worker_task:
            self._worker_task.cancel()
            try:
                await self._worker_task
            except asyncio.CancelledError:
                pass
        print("[TaskManager] 已停止")

    async def submit_task(self, description: str, project_path: str,
                          requirements: str = "", extra: dict = None) -> dict:
        """提交一个新任务"""
        task_id = str(uuid.uuid4())[:8]
        now = time.time()
        task = {
            "id": task_id,
            "description": description,
            "project_path": project_path,
            "requirements": requirements,
            "status": STATUS_PENDING,
            "created_at": now,
            "updated_at": now,
            "result": None,
            "logs": [],
            "extra": extra or {},
        }
        self._tasks[task_id] = task
        self.store.save(task)
        await self._queue.put(task_id)
        print(f"[TaskManager] 任务已提交: {task_id} - {description[:50]}")
        return task

    def get_task(self, task_id: str) -> Optional[dict]:
        """获取任务详情"""
        # 优先从内存取，没有再从文件读
        if task_id in self._tasks:
            return self._tasks[task_id]
        return self.store.load(task_id)

    def get_task_status(self, task_id: str) -> Optional[str]:
        task = self.get_task(task_id)
        return task["status"] if task else None

    def list_tasks(self, status: str = None) -> List[dict]:
        """列出所有任务，可按状态过滤"""
        tasks = list(self._tasks.values())
        # 加上磁盘上的任务
        disk_tasks = self.store.list_all()
        disk_ids = {t["id"] for t in tasks}
        for t in disk_tasks:
            if t["id"] not in disk_ids:
                tasks.append(t)
        if status:
            tasks = [t for t in tasks if t["status"] == status]
        tasks.sort(key=lambda t: t.get("created_at", 0), reverse=True)
        return tasks

    def append_log(self, task_id: str, message: str):
        """追加任务日志"""
        task = self._tasks.get(task_id)
        if not task:
            return
        if "logs" not in task:
            task["logs"] = []
        task["logs"].append({
            "time": time.strftime("%Y-%m-%d %H:%M:%S"),
            "message": message,
        })
        if len(task["logs"]) > 500:
            task["logs"] = task["logs"][-500:]
        # 异步持久化（这里直接同步写，简单起见）
        self.store.save(task)

    async def cancel_task(self, task_id: str) -> bool:
        """取消任务"""
        task = self._tasks.get(task_id)
        if not task:
            return False
        if task["status"] in (STATUS_COMPLETED, STATUS_FAILED, STATUS_CANCELLED):
            return False
        task["status"] = STATUS_CANCELLED
        task["updated_at"] = time.time()
        self.store.save(task)
        print(f"[TaskManager] 任务已取消: {task_id}")
        return True

    async def _worker_loop(self):
        """后台 worker，从队列取任务执行"""
        while True:
            try:
                task_id = await self._queue.get()
                async with self._semaphore:
                    await self._execute_task(task_id)
            except asyncio.CancelledError:
                break
            except Exception as e:
                print(f"[TaskManager] Worker 异常: {e}")
                await asyncio.sleep(1)

    async def _execute_task(self, task_id: str):
        """执行单个任务"""
        task = self._tasks.get(task_id)
        if not task or task["status"] == STATUS_CANCELLED:
            return

        task["status"] = STATUS_RUNNING
        task["updated_at"] = time.time()
        self.store.save(task)
        print(f"[TaskManager] 开始执行: {task_id}")

        try:
            result = await self._executor(task)
            task["result"] = result
            # 检查是否已被取消
            if task["status"] == STATUS_CANCELLED:
                return
            task["status"] = STATUS_COMPLETED
            print(f"[TaskManager] 任务完成: {task_id}")
        except Exception as e:
            task["status"] = STATUS_FAILED
            task["result"] = {"error": str(e)}
            self.append_log(task_id, f"执行失败: {e}")
            print(f"[TaskManager] 任务失败: {task_id} - {e}")
        finally:
            task["updated_at"] = time.time()
            self.store.save(task)