"""Background worker for queued transcription jobs."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from contextlib import contextmanager
from contextlib import suppress
from datetime import UTC
from datetime import datetime
from uuid import UUID

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

from transcription.config import Settings
from transcription.config import get_settings
from transcription.db import get_session
from transcription.errors import AppError
from transcription.errors import ErrorCategory
from transcription.errors import classify_unexpected_error
from transcription.errors import format_error_detail
from transcription.models import Job
from transcription.models import JobStatus
from transcription.models import Transcript
from transcription.services.transcription import TranscriptionError
from transcription.services.transcription import transcribe_document_image

from .services import ServiceBundle
from .services.jobs import JobService

logger = logging.getLogger(__name__)


async def queue_consumer_loop(queue: asyncio.Queue[UUID], stop_event: asyncio.Event):
    """Main worker loop that consumes jobs from the queue and processes them.

    The queue is for Job UUIDs, and the corresponding documents should already have been uploaded.
    """
    service = JobService()
    while not stop_event.is_set():
        with handle_worker_exceptions():
            async with _get_queue_item(queue) as job_id:
                job = await service.read_job(job_id)
                asyncio.create_task(process_job(job=job, services=ServiceBundle()))


@contextmanager
def handle_worker_exceptions():
    """Context manager to log and suppress exceptions in the worker loop."""
    try:
        yield
    except Exception as exc:
        error = exc if isinstance(exc, AppError) else classify_unexpected_error(exc, operation="worker.loop")
        logger.exception(
            "Worker loop exception error_id=%s category=%s",
            error.error_id,
            error.category.value,
        )


@asynccontextmanager
async def _get_queue_item(queue: asyncio.Queue[UUID]) -> AsyncGenerator[UUID]:
    """Context manager to enqueue a job and ensure it is marked done."""
    yield await queue.get()
    queue.task_done()


async def process_job(job: Job, services: ServiceBundle) -> None:
    """Process a single job using the service bundle."""
    match job.status:
        case JobStatus.QUEUED:
            job.status = JobStatus.PROCESSING
            try:
                await services.transcriptions.transcribe_image(job.document.file_path)
                job.status = JobStatus.TRANSCRIBED
            except TranscriptionError as exc:
                job.status = JobStatus.FAILED
                job.error_message = str(exc)
            await services.jobs.update_job(job)


async def run_worker_loop(
    *,
    session_factory: async_sessionmaker[AsyncSession] | None = None,
    stop_event: asyncio.Event | None = None,
    wake_queue: asyncio.Queue[None] | None = None,
    poll_interval_seconds: float = 1.0,
) -> None:
    """Run worker loop until stop_event is set.

    If wake_queue is provided, queue activity wakes the loop immediately while
    timeout-based wakeups preserve current polling behavior.
    """
    while True:
        if stop_event is not None and stop_event.is_set():
            logger.info("Worker stop event received")
            return

        if wake_queue is not None:
            with suppress(TimeoutError):
                await asyncio.wait_for(wake_queue.get(), timeout=poll_interval_seconds)

        processed_any = False
        while await process_next_queued_job(session_factory=session_factory):
            processed_any = True

        if wake_queue is None and not processed_any:
            await asyncio.sleep(poll_interval_seconds)


async def process_next_queued_job(
    *,
    session: AsyncSession | None = None,
    session_factory: async_sessionmaker[AsyncSession] | None = None,
) -> bool:
    """Process the next queued job and persist terminal outcome.

    Returns True when a job was processed, False when no queued job exists.
    """
    if session is None:
        async with get_session(session_factory=session_factory) as local_session:
            return await _process_next_queued_job(session=local_session)
    return await _process_next_queued_job(session=session)


async def _process_next_queued_job(*, session: AsyncSession) -> bool:
    query = (
        select(Job)
        .options(selectinload(Job.document))  # pyright: ignore[reportArgumentType]
        .where(Job.status == JobStatus.QUEUED)
        .order_by(Job.created_at)  # pyright: ignore[reportArgumentType]
    )
    job = (await session.exec(query)).first()

    if job is None:
        return False

    logger.info("Picked queued job operation=worker.pick job_id=%s", job.id)
    job.status = JobStatus.PROCESSING
    job.updated_at = datetime.now(UTC)
    session.add(job)
    await session.commit()
    await session.refresh(job)

    document = job.document
    if document is None:
        error = AppError(
            "Document not found",
            category=ErrorCategory.NOT_FOUND,
            suggestion="Re-upload the source document and retry processing.",
        )
        await _finalize_failed_job(session=session, job=job, error=error)
        logger.error(
            "Job failed operation=worker.process_job job_id=%s error_id=%s category=%s",
            job.id,
            error.error_id,
            error.category.value,
        )
        return True

    try:
        # Provider SDK calls are synchronous and should not block the event loop.
        result = await asyncio.to_thread(transcribe_document_image, document.file_path)
        await _upsert_transcript(session=session, job_id=job.id, text=result.text, error_detail=None)
        job.status = JobStatus.TRANSCRIBED
        job.updated_at = datetime.now(UTC)
        session.add(job)
        await session.commit()
        logger.info(
            "Job transcribed operation=worker.process_job job_id=%s document_id=%s provider=%s",
            job.id,
            document.id,
            result.provider,
        )
    except Exception as exc:
        error = exc if isinstance(exc, AppError) else classify_unexpected_error(exc, operation="worker.process_job")
        settings = get_settings()
        if _should_retry(job=job, error=error, settings=settings):
            await _requeue_for_retry(session=session, job=job, error=error, settings=settings)
            logger.warning(
                "Job retried operation=worker.process_job "
                "job_id=%s document_id=%s retry_count=%s error_id=%s category=%s",
                job.id,
                document.id,
                job.retry_count,
                error.error_id,
                error.category.value,
            )
        else:
            await _finalize_failed_job(session=session, job=job, error=error)
            logger.exception(
                "Job failed operation=worker.process_job job_id=%s document_id=%s error_id=%s category=%s",
                job.id,
                document.id,
                error.error_id,
                error.category.value,
            )

    return True


async def _upsert_transcript(
    *, session: AsyncSession, job_id, text: str | None, error_detail: str | None
) -> Transcript:
    transcript = (await session.exec(select(Transcript).where(Transcript.job_id == job_id))).first()
    if transcript is None:
        transcript = Transcript(job_id=job_id)

    transcript.text = text
    transcript.error_detail = error_detail
    session.add(transcript)
    await session.commit()
    await session.refresh(transcript)
    return transcript


def _should_retry(*, job: Job, error: AppError, settings: Settings) -> bool:
    return error.retriable and job.retry_count < settings.worker_max_retries


async def _requeue_for_retry(*, session: AsyncSession, job: Job, error: AppError, settings: Settings) -> None:
    await _upsert_transcript(session=session, job_id=job.id, text=None, error_detail=format_error_detail(error))
    job.retry_count += 1
    job.status = JobStatus.QUEUED
    job.updated_at = datetime.now(UTC)
    session.add(job)
    await session.commit()
    if settings.worker_retry_backoff_seconds > 0:
        await asyncio.sleep(settings.worker_retry_backoff_seconds)


async def _finalize_failed_job(*, session: AsyncSession, job: Job, error: AppError) -> None:
    await _upsert_transcript(session=session, job_id=job.id, text=None, error_detail=format_error_detail(error))
    job.status = JobStatus.FAILED
    job.updated_at = datetime.now(UTC)
    session.add(job)
    await session.commit()
