"""
任务调度器模块

协调任务队列、执行器和通知器，管理任务生命周期。
"""

import asyncio
import logging
from typing import Dict, Any, Optional, List

from src.services.task.queue import TaskQueue, TaskPriority
from src.services.task.executor import TaskExecutor, ExecutionStatus, TaskExecutionResult
from src.repositories.task_repository_async import AsyncTaskRepository
from src.config.database_async import get_async_session_context
from src.models.entities import TaskDB

logger = logging.getLogger(__name__)


class TaskScheduler:
    """
    任务调度器

    协调以下组件:
    - TaskQueue: 管理任务队列
    - TaskExecutor: 执行任务
    - AsyncTaskRepository: 持久化任务状态

    特性:
    - 工作线程池管理
    - 任务生命周期管理
    - 自动恢复
    - 优雅关闭
    """

    def __init__(
        self,
        num_workers: int = 5,
        queue_maxsize: int = 1000
    ):
        """
        初始化任务调度器

        Args:
            num_workers: 工作线程数量
            queue_maxsize: 队列最大容量
        """
        # 组件
        self._queue = TaskQueue(maxsize=queue_maxsize)
        self._executor = TaskExecutor()

        # 工作线程
        self._num_workers = num_workers
        self._workers: List[asyncio.Task] = []
        self._running = False
        self._shutdown_event = asyncio.Event()

        # 取消事件映射
        self._cancellation_events: Dict[str, asyncio.Event] = {}
        self._cancellation_lock = asyncio.Lock()

        # 统计
        self._stats = {
            "tasks_submitted": 0,
            "tasks_completed": 0,
            "tasks_failed": 0,
            "tasks_retried": 0,
            "tasks_cancelled": 0
        }

        logger.info(f"TaskScheduler initialized with {num_workers} workers")

    @property
    def queue(self) -> TaskQueue:
        """获取任务队列"""
        return self._queue

    @property
    def executor(self) -> TaskExecutor:
        """获取任务执行器"""
        return self._executor


    async def start(self):
        """启动调度器"""
        if self._running:
            logger.warning("TaskScheduler is already running")
            return

        self._running = True
        self._shutdown_event.clear()

        # 恢复未完成的任务
        await self._recover_pending_tasks()

        # 启动工作线程
        for i in range(self._num_workers):
            worker = asyncio.create_task(
                self._worker_loop(f"worker-{i}"),
                name=f"task-worker-{i}"
            )
            self._workers.append(worker)

        logger.info(f"TaskScheduler started with {self._num_workers} workers")

    async def stop(self, wait_for_queue: bool = True, timeout: float = 60.0):
        """
        停止调度器

        Args:
            wait_for_queue: 是否等待队列中的任务完成
            timeout: 等待超时时间
        """
        if not self._running:
            return

        logger.info("TaskScheduler stopping...")
        self._running = False
        self._shutdown_event.set()

        if wait_for_queue and not self._queue.is_empty:
            try:
                await asyncio.wait_for(self._queue.wait_for_empty(), timeout=timeout)
            except asyncio.TimeoutError:
                logger.warning("Timeout waiting for queue to empty")

        # 取消所有工作线程
        for worker in self._workers:
            worker.cancel()

        # 等待工作线程结束
        await asyncio.gather(*self._workers, return_exceptions=True)
        self._workers.clear()

        logger.info("TaskScheduler stopped")

    async def submit_task(
        self,
        task_id: str,
        task_type: str,
        provider: str,
        model: str,
        params: Dict[str, Any],
        user_id: Optional[str] = None,
        project_id: Optional[str] = None,
        priority: TaskPriority = TaskPriority.NORMAL
    ) -> bool:
        """
        提交任务到调度器

        Args:
            task_id: 任务 ID
            task_type: 任务类型
            provider: 提供商
            model: 模型
            params: 任务参数
            user_id: 用户 ID
            project_id: 项目 ID
            priority: 优先级

        Returns:
            是否成功提交
        """
        try:
            # 加入队列
            success = await self._queue.enqueue(
                task_id=task_id,
                task_type=task_type,
                priority=priority,
                user_id=user_id,
                project_id=project_id
            )

            if success:
                self._stats["tasks_submitted"] += 1

                logger.info(f"Task submitted: {task_id} (type={task_type}, priority={priority.name})")

            return success

        except Exception as e:
            logger.error(f"Failed to submit task {task_id}: {e}")
            return False

    async def cancel_task(self, task_id: str) -> bool:
        """
        取消任务

        Args:
            task_id: 任务 ID

        Returns:
            是否成功取消
        """
        # 设置取消事件
        async with self._cancellation_lock:
            if task_id in self._cancellation_events:
                self._cancellation_events[task_id].set()

        # 尝试从队列中移除（如果还在队列中）
        # 注意：TaskQueue 不支持直接移除，任务会在出队时检查取消状态

        # 尝试取消正在执行的任务
        cancelled = await self._executor.cancel_task(task_id)

        if cancelled:
            self._stats["tasks_cancelled"] += 1

            # 更新数据库状态
            async with get_async_session_context() as session:
                repo = AsyncTaskRepository(session)
                await repo.update_task_status(
                    task_id=task_id,
                    status="cancelled"
                )

            logger.info(f"Task cancelled: {task_id}")

        return cancelled

    async def _worker_loop(self, worker_name: str):
        """
        工作线程主循环

        Args:
            worker_name: 工作线程名称
        """
        logger.info(f"Worker {worker_name} started")

        while self._running:
            try:
                # 从队列获取任务
                item = await self._queue.dequeue(timeout=1.0)

                if item is None:
                    continue

                # 检查取消状态
                async with self._cancellation_lock:
                    cancellation_event = self._cancellation_events.get(item.task_id)
                    if cancellation_event and cancellation_event.is_set():
                        logger.debug(f"Task {item.task_id} was cancelled, skipping")
                        continue

                    # 创建新的取消事件
                    cancellation_event = asyncio.Event()
                    self._cancellation_events[item.task_id] = cancellation_event

                # 执行任务
                await self._execute_task(item, cancellation_event)

                # 清理取消事件
                async with self._cancellation_lock:
                    self._cancellation_events.pop(item.task_id, None)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.exception(f"Worker {worker_name} error: {e}")

        logger.info(f"Worker {worker_name} stopped")

    async def _execute_task(self, item, cancellation_event: asyncio.Event):
        """
        执行单个任务

        Args:
            item: 任务队列项
            cancellation_event: 取消事件
        """
        task_id = item.task_id
        task_type = item.task_type

        logger.debug(f"Executing task: {task_id}")

        # 获取任务信息
        async with get_async_session_context() as session:
            repo = AsyncTaskRepository(session)
            task_db = await repo.get(task_id)

            if not task_db:
                logger.error(f"Task not found in database: {task_id}")
                return

            # 更新状态为处理中
            await repo.update_task_status(
                task_id=task_id,
                status="processing"
            )

        # 提取提供商和模型
        provider = task_db.provider or "default"
        model = task_db.model or "default"
        params = task_db.params or {}

        # 执行
        result = await self._executor.execute(
            task_id=task_id,
            task_type=task_type,
            provider=provider,
            model=model,
            params=params,
            retry_count=task_db.retry_count,
            cancellation_event=cancellation_event
        )

        # 处理结果
        await self._handle_execution_result(task_id, item, result)

    async def _handle_execution_result(
        self,
        task_id: str,
        item,
        result: TaskExecutionResult
    ):
        """
        处理执行结果

        Args:
            task_id: 任务 ID
            item: 任务队列项
            result: 执行结果
        """
        task_type = item.task_type

        async with get_async_session_context() as session:
            repo = AsyncTaskRepository(session)

            if result.status == ExecutionStatus.SUCCESS:
                # 成功
                await repo.update_task_status(
                    task_id=task_id,
                    status="success",
                    result=result.result
                )

                self._stats["tasks_completed"] += 1

                logger.info(f"Task completed: {task_id}")

            elif result.status == ExecutionStatus.CANCELLED:
                # 取消
                await repo.update_task_status(
                    task_id=task_id,
                    status="cancelled"
                )

                self._stats["tasks_cancelled"] += 1

                logger.info(f"Task cancelled: {task_id}")

            elif result.status == ExecutionStatus.RETRYABLE:
                # 可重试
                should_retry = await self._executor.should_retry(
                    task_type=task_type,
                    retry_count=item.retry_count or 0,
                    result=result
                )

                if should_retry:
                    # 计算重试延迟
                    delay = self._executor.calculate_retry_delay(
                        task_type=task_type,
                        retry_count=item.retry_count or 0
                    )

                    # 更新重试状态
                    await repo.update_by_id(task_id, {
                        "status": "retrying",
                        "retry_count": (item.retry_count or 0) + 1,
                        "error": result.error
                    })

                    self._stats["tasks_retried"] += 1

                    # 重新加入队列
                    await asyncio.sleep(delay)
                    await self._queue.enqueue(
                        task_id=task_id,
                        task_type=task_type,
                        priority=TaskPriority.NORMAL,
                        user_id=item.user_id,
                        project_id=item.project_id
                    )

                    logger.info(f"Task retrying: {task_id} (delay={delay:.1f}s)")

                else:
                    # 重试次数耗尽，标记为失败
                    await repo.update_task_status(
                        task_id=task_id,
                        status="failed",
                        error=result.error or "Max retries exceeded"
                    )

                    self._stats["tasks_failed"] += 1

                    logger.error(f"Task failed after retries: {task_id}")

            else:
                # 失败（不可重试）
                await repo.update_task_status(
                    task_id=task_id,
                    status="failed",
                    error=result.error
                )

                self._stats["tasks_failed"] += 1

                logger.error(f"Task failed: {task_id} - {result.error}")

    async def _recover_pending_tasks(self):
        """
        恢复未完成的任务

        在系统重启后，将所有未完成的任务重新加入队列。
        """
        try:
            async with get_async_session_context() as session:
                repo = AsyncTaskRepository(session)

                # 查询所有未完成的任务
                pending_tasks = await repo.list_active_tasks(limit=1000)

                recovered_count = 0
                for task_db in pending_tasks:
                    # 重置状态为 pending
                    await repo.update_task_status(
                        task_id=task_db.id,
                        status="pending"
                    )

                    # 加入队列
                    await self._queue.enqueue(
                        task_id=task_db.id,
                        task_type=task_db.type,
                        priority=TaskPriority.NORMAL,
                        user_id=task_db.user_id,
                        project_id=task_db.project_id
                    )

                    recovered_count += 1

                if recovered_count > 0:
                    logger.info(f"Recovered {recovered_count} pending tasks")

        except Exception as e:
            logger.error(f"Failed to recover pending tasks: {e}")

    def get_stats(self) -> Dict[str, Any]:
        """获取调度器统计信息"""
        return {
            **self._stats,
            "queue_stats": self._queue.get_stats(),
            "executor_stats": self._executor.get_stats(),
            "running": self._running,
            "num_workers": self._num_workers
        }
