"""Batch queue worker."""

from __future__ import annotations

import asyncio
import logging

from sqlalchemy import select

from app.db import get_session_factory
from app.models import Batch, Document
from app.schemas import DocStatus
from app.services.pipeline import process_document
from app.timeutil import utc_now

logger = logging.getLogger(__name__)

_running: set[str] = set()


async def run_batch(batch_id: str) -> None:
    """Process all queued documents in a batch."""
    if batch_id in _running:
        return
    _running.add(batch_id)
    try:
        sessions = get_session_factory()
        async with sessions() as session:
            batch = await session.get(Batch, batch_id)
            if not batch:
                return
            batch.status = "processing"
            await session.commit()

            result = await session.execute(
                select(Document).where(
                    Document.batch_id == batch_id,
                    Document.status == DocStatus.QUEUED.value,
                )
            )
            docs = list(result.scalars().all())

            for doc in docs:
                await process_document(session, doc)
                batch = await session.get(Batch, batch_id)
                if not batch:
                    return
                batch.processed += 1
                if doc.status == DocStatus.FAILED.value:
                    batch.failed += 1
                await session.commit()

            batch = await session.get(Batch, batch_id)
            if batch:
                batch.status = "completed"
                batch.finished_at = utc_now()
                await session.commit()
    except Exception as exc:  # noqa: BLE001
        logger.exception("batch %s failed: %s", batch_id, exc)
        try:
            sessions = get_session_factory()
        except RuntimeError:
            return
        async with sessions() as session:
            batch = await session.get(Batch, batch_id)
            if batch:
                batch.status = "failed"
                batch.finished_at = utc_now()
                await session.commit()
    finally:
        _running.discard(batch_id)


def enqueue_batch(batch_id: str) -> None:
    """Fire-and-forget batch processing task."""
    asyncio.create_task(run_batch(batch_id))