"""
异步任务仓库

使用 BaseRepository 模式提供任务的异步数据访问操作。
"""

from typing import Optional, List
from sqlalchemy.ext.asyncio import AsyncSession
from src.models.entities import TaskDB
from src.repositories.base_async import AsyncBaseRepository
import logging

logger = logging.getLogger(__name__)


class AsyncTaskRepository(AsyncBaseRepository[TaskDB]):
    """
    Task 实体的异步仓库。

    继承自 AsyncBaseRepository 的所有异步 CRUD 操作，并添加任务特定的异步查询。
    """

    def __init__(self, session: AsyncSession):
        """
        初始化 AsyncTaskRepository。

        Args:
            session: 异步数据库会话
        """
        super().__init__(TaskDB, session)

    async def get_by_provider_task_id(self, provider_task_id: str) -> Optional[TaskDB]:
        """
        通过提供商任务 ID 获取任务。

        Args:
            provider_task_id: AI 提供商的任务 ID

        Returns:
            Task 实例或 None
        """
        return await self.get_by_field("provider_task_id", provider_task_id)

    async def list_by_status(
        self,
        status: str,
        limit: int = 100,
        offset: int = 0
    ) -> List[TaskDB]:
        """
        按状态列出任务。

        Args:
            status: 要过滤的任务状态
            limit: 返回的最大任务数
            offset: 要跳过的任务数

        Returns:
            任务列表
        """
        return await self.list(
            skip=offset,
            limit=limit,
            filters={"status": status, "deleted_at": None},
            sort_by="created_at",
            sort_order="desc"
        )

    async def list_by_user(
        self,
        user_id: str,
        limit: int = 100,
        offset: int = 0
    ) -> List[TaskDB]:
        """
        列出特定用户的任务。

        Args:
            user_id: 要过滤的用户 ID
            limit: 返回的最大任务数
            offset: 要跳过的任务数

        Returns:
            任务列表
        """
        return await self.list(
            skip=offset,
            limit=limit,
            filters={"user_id": user_id, "deleted_at": None},
            sort_by="created_at",
            sort_order="desc"
        )

    async def list_by_project(
        self,
        project_id: str,
        limit: int = 100,
        offset: int = 0
    ) -> List[TaskDB]:
        """
        列出特定项目的任务。

        Args:
            project_id: 要过滤的项目 ID
            limit: 返回的最大任务数
            offset: 要跳过的任务数

        Returns:
            任务列表
        """
        return await self.list(
            skip=offset,
            limit=limit,
            filters={"project_id": project_id, "deleted_at": None},
            sort_by="created_at",
            sort_order="desc"
        )

    async def list_pending_tasks(self, limit: int = 100) -> List[TaskDB]:
        """
        列出所有待处理的任务。

        Args:
            limit: 返回的最大任务数

        Returns:
            待处理任务列表
        """
        return await self.list(
            limit=limit,
            filters={"status": "pending", "deleted_at": None},
            sort_by="created_at",
            sort_order="asc"
        )

    async def list_active_tasks(self, limit: int = 100) -> List[TaskDB]:
        """
        列出所有活动中的任务（待处理、处理中、重试中）。

        Args:
            limit: 返回的最大任务数

        Returns:
            活动任务列表
        """
        return await self.list(
            limit=limit,
            filters={"status__in": ["pending", "processing", "retrying"], "deleted_at": None},
            sort_by="created_at",
            sort_order="asc"
        )

    async def list_stuck_tasks(self, timeout_seconds: int = 3600) -> List[TaskDB]:
        """
        列出卡在处理状态的任务。

        Args:
            timeout_seconds: 任务被视为卡住之前的秒数

        Returns:
            卡住的任务列表
        """
        from datetime import datetime
        cutoff_time = datetime.now().timestamp() - timeout_seconds

        return await self.list(
            filters={
                "status__in": ["processing", "retrying"],
                "updated_at__lt": cutoff_time,
                "deleted_at": None
            },
            sort_by="updated_at",
            sort_order="asc"
        )

    async def list_by_type(self, task_type: str, limit: int = 100, offset: int = 0) -> List[TaskDB]:
        """
        按类型列出任务。

        Args:
            task_type: 任务类型 (image, video, audio, music, script)
            limit: 返回的最大任务数
            offset: 要跳过的任务数

        Returns:
            任务列表
        """
        return await self.list(
            skip=offset,
            limit=limit,
            filters={"type": task_type, "deleted_at": None},
            sort_by="created_at",
            sort_order="desc"
        )

    async def count_by_status(self, status: str) -> int:
        """
        按状态计数任务。

        Args:
            status: 要计数的任务状态

        Returns:
            具有给定状态的任务数
        """
        return await self.count(filters={"status": status, "deleted_at": None})

    async def count_by_user(self, user_id: str) -> int:
        """
        计数特定用户的任务。

        Args:
            user_id: 要计数任务的用户 ID

        Returns:
            该用户的任务数
        """
        return await self.count(filters={"user_id": user_id, "deleted_at": None})

    async def count_by_project(self, project_id: str) -> int:
        """
        计数特定项目的任务。

        Args:
            project_id: 要计数任务的项目 ID

        Returns:
            该项目的任务数
        """
        return await self.count(filters={"project_id": project_id, "deleted_at": None})

    async def update_task_status(
        self,
        task_id: str,
        status: str,
        result: Optional[dict] = None,
        error: Optional[str] = None
    ) -> Optional[TaskDB]:
        """
        更新任务状态。

        Args:
            task_id: 任务 ID
            status: 新状态
            result: 可选的结果数据
            error: 可选的错误信息

        Returns:
            更新后的任务或 None
        """
        from datetime import datetime

        update_data = {
            "status": status,
            "updated_at": datetime.now().timestamp()
        }

        if result is not None:
            update_data["result"] = result

        if error is not None:
            update_data["error"] = error

        if status in ["success", "failed", "cancelled", "timeout"]:
            update_data["completed_at"] = datetime.now().timestamp()

        return await self.update_by_id(task_id, update_data)
