from __future__ import annotations

from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from models.task import Task


class TaskRepository:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def list(self, limit: int = 100, offset: int = 0) -> list[Task]:
        stmt = select(Task).order_by(Task.created_at.desc()).offset(offset).limit(limit)
        result = await self.session.execute(stmt)
        return result.scalars().all()

    async def get(self, task_id: str) -> Optional[Task]:
        stmt = select(Task).where(Task.id == task_id).options(selectinload(Task.schedule_runs))
        result = await self.session.execute(stmt)
        return result.scalar_one_or_none()

    async def create(self, *, id: str, action: str, target: str, playbook: Optional[str] = None,
                     status: str = "pending") -> Task:
        task = Task(
            id=id,
            action=action,
            target=target,
            playbook=playbook,
            status=status,
        )
        self.session.add(task)
        await self.session.flush()
        return task

    async def update(self, task: Task, **fields) -> Task:
        for key, value in fields.items():
            if value is not None:
                setattr(task, key, value)
        await self.session.flush()
        return task