from sqlalchemy import select
from sqlalchemy.orm import selectinload
from src.db.models import tasks as task_models
from src.schemas.tasks import task as task_schemas
from src.schemas.tasks import runtime as task_runtime_schemas
from src.utils.retention import RetentionOption, get_retention_cutoff
from .resource import TaskResourceService
from ..service_base import ServiceBase
from ..exceptions import NotFoundError, ServiceErrorCode


class TaskNotFoundError(NotFoundError):
    def __init__(self, task_id: int) -> None:
        super().__init__(ServiceErrorCode.TASK_NOT_FOUND, "Task", task_id)

class TaskService(ServiceBase[task_models.Task]):
    @staticmethod
    def relations():
        return [
            selectinload(task_models.Task.agent),
            selectinload(task_models.Task.workspace),
        ]

    def get_tasks_query(self, workspace_id: int):
        return (
            select(task_models.Task)
            .where(task_models.Task.workspace_id == workspace_id)
            .order_by(task_models.Task.id.desc())
        )

    def get_recent_tasks_query(self):
        return (
            select(task_models.Task)
            .order_by(
                task_models.Task.last_run_at.desc(),
                task_models.Task.id.desc(),
            )
        )

    async def get_task_by_id(self, id: int) -> task_models.Task:
        task = await self._db_session.get(
            task_models.Task, id,
            options=self.relations(),
        )
        if not task:
            raise TaskNotFoundError(id)
        return task

    async def create_task(self, data: task_schemas.TaskCreate) -> task_models.Task:
        new_task = task_models.Task(
            _workspace_id=data.workspace_id,
            **data.model_dump(exclude={"workspace_id"})
        )

        self._db_session.add(new_task)
        new_id = await self.flush_and_expunge(new_task)
        return await self.get_task_by_id(new_id)

    async def update_task(self, id: int, data: task_schemas.TaskUpdate) -> task_models.Task:
        task = await self.get_task_by_id(id)

        if data.messages is not None:
            task.messages = data.messages

        self.apply_fields(task, data, exclude={"messages"})

        new_id = await self.flush_and_expunge(task)
        return await self.get_task_by_id(new_id)

    async def delete_task(self, id: int):
        task = await self.get_task_by_id(id)
        await self._db_session.delete(task)
        await self._db_session.flush()
        await TaskResourceService(self._db_session, task_runtime_schemas.TaskType.TASK).delete_task_resources(id)

    async def cleanup_outdated_tasks(self, retention: RetentionOption):
        cutoff = get_retention_cutoff(retention)
        if cutoff is None: return
        stmt = select(task_models.Task.id).where(task_models.Task.last_run_at < cutoff)
        task_ids = (await self._db_session.scalars(stmt)).all()

        for task_id in task_ids:
            await self.delete_task(task_id)
