from __future__ import annotations

from datetime import datetime

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

from app.models.entities import Job


class JobRepository:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def create(self, job: Job) -> Job:
        self.session.add(job)
        await self.session.flush()
        await self.session.refresh(job)
        return job

    async def get_by_id(self, job_id: str) -> Job | None:
        return await self.session.get(Job, job_id)

    async def list_all(self) -> list[Job]:
        result = await self.session.execute(select(Job).order_by(Job.created_at.desc()))
        return list(result.scalars().all())

    async def list_runnable(self, now: datetime, limit: int) -> list[Job]:
        result = await self.session.execute(
            select(Job)
            .where(
                Job.status.in_(["queued", "retryable"]),
                Job.run_after <= now,
            )
            .order_by(Job.run_after.asc(), Job.created_at.asc())
            .limit(limit)
        )
        return list(result.scalars().all())

    async def find_active_by_kind_and_payload(self, kind: str, payload_json: str) -> Job | None:
        result = await self.session.execute(
            select(Job).where(
                Job.kind == kind,
                Job.payload_json == payload_json,
                Job.status.in_(["queued", "retryable", "running"]),
            )
        )
        return result.scalar_one_or_none()

    async def save(self, job: Job) -> Job:
        await self.session.flush()
        await self.session.refresh(job)
        return job