from collections.abc import Sequence
from datetime import UTC
from datetime import datetime
from uuid import UUID

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

from ..models import Job
from ..models import JobStatus
from .base import ServiceBase


class JobService(ServiceBase):
    """Thin service class for managing jobs in the database."""

    #
    # CRUD Operations
    #

    async def create_job(self, job: Job, session: AsyncSession | None = None) -> Job:
        """Create a new job in the database."""
        async with self._session_scope(session) as _session:
            _session.add(job)
            await self._finalize(session=_session, caller_session=session, refresh=(job,))
            return job

    async def read_job(self, job_id: UUID, session: AsyncSession | None = None) -> Job:
        """Read an existing job from the database.

        The related document is always eagerly loaded so callers can safely
        access ``job.document`` in async contexts without triggering lazy-load IO.
        """
        async with self._session_scope(session) as _session:
            query = (
                select(Job)
                .options(
                    selectinload(Job.document),  # pyright: ignore[reportArgumentType]
                    selectinload(Job.transcripts),  # pyright: ignore[reportArgumentType]
                )
                .where(Job.id == job_id)
                .execution_options(populate_existing=True)
            )
            job = (await _session.exec(query)).first()
        if job is None:
            raise ValueError(f"Job with id {job_id} not found")
        return job

    async def update_job(self, job: Job, session: AsyncSession | None = None) -> Job:
        """Update an existing job in the database."""
        async with self._session_scope(session) as _session:
            merged = await _session.merge(job)
            await self._finalize(session=_session, caller_session=session, refresh=(merged,))
            return merged

    async def delete_job(self, job: Job, session: AsyncSession | None = None) -> None:
        """Delete a job from the database."""
        async with self._session_scope(session) as _session:
            await _session.delete(job)
            await self._finalize(session=_session, caller_session=session)

    # Query Operations

    async def query_jobs(
        self,
        *,
        status: JobStatus | None = None,
        filename: str | None = None,
        session: AsyncSession | None = None,
    ) -> Sequence[Job]:
        """Query jobs from the database based on provided filters."""
        async with self._session_scope(session) as _session:
            query = select(Job).options(selectinload(Job.document))  # pyright: ignore[reportArgumentType]
            if status is not None:
                query = query.where(Job.status == status)
            if filename is not None:
                query = query.where(Job.document.filename == filename)
            result = await _session.exec(query)
            return result.all()

    async def list_jobs(
        self,
        *,
        load_docs: bool = False,
        session: AsyncSession | None = None,
    ) -> Sequence[Job]:
        """List all jobs in the database with eagerly loaded documents."""
        _ = load_docs
        async with self._session_scope(session) as _session:
            query = select(Job).options(selectinload(Job.document))  # pyright: ignore[reportArgumentType]
            result = await _session.exec(query)
            return result.all()

    # Other Operations

    async def mark_job_status(
        self,
        job_id: UUID,
        status: JobStatus,
        session: AsyncSession | None = None,
    ) -> Job:
        """Mark a job with a new status."""
        return await self.update_job_state(job_id=job_id, status=status, session=session)

    async def update_job_state(
        self,
        *,
        job_id: UUID,
        status: JobStatus,
        retry_count_increment: int = 0,
        session: AsyncSession | None = None,
    ) -> Job:
        """Update a job's lifecycle fields.

        When ``session`` is provided, this method flushes so callers can commit
        once at an orchestration boundary.
        """
        async with self._session_scope(session) as _session:
            query = (
                select(Job)
                .options(selectinload(Job.document))  # pyright: ignore[reportArgumentType]
                .where(Job.id == job_id)
                .execution_options(populate_existing=True)
            )
            job = (await _session.exec(query)).first()
            if job is None:
                raise ValueError(f"Job with id {job_id} not found")
            job.status = status
            if retry_count_increment:
                job.retry_count += retry_count_increment
            job.updated_at = datetime.now(UTC)
            await self._finalize(session=_session, caller_session=session, refresh=(job,))
            return job

    async def read_next_queued_job(
        self,
        *,
        session: AsyncSession | None = None,
    ) -> Job | None:
        """Read the next queued job ordered by creation time."""
        async with self._session_scope(session) as _session:
            query = (
                select(Job)
                .options(selectinload(Job.document))  # pyright: ignore[reportArgumentType]
                .where(Job.status == JobStatus.QUEUED)
                .order_by(Job.created_at)  # pyright: ignore[reportArgumentType]
            )
            return (await _session.exec(query)).first()
