UmbraAudio/umbraaudio/cluster/jobs.py

"""Job queue operations for UmbraAudio worker."""

from __future__ import annotations

import asyncio
import contextlib
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from .database import ClusterDB, JobRecord, get_cluster_db

if TYPE_CHECKING:
    import uuid

    from .config import ClusterConfig

logger = logging.getLogger(__name__)


class JobKind:
    """Job kinds that UmbraAudio can claim."""
    TTS = "tts"
    ASR = "asr"
    ASSEMBLE = "assemble"
    ENCODE = "encode"
    MASTER = "master"
    SPEAKER_PROFILE = "speaker_profile"

    @classmethod
    def all(cls) -> list[str]:
        return [cls.TTS, cls.ASR, cls.ASSEMBLE, cls.ENCODE, cls.MASTER, cls.SPEAKER_PROFILE]

    @classmethod
    def tts_related(cls) -> list[str]:
        return [cls.TTS, cls.ASSEMBLE, cls.ENCODE, cls.MASTER]


@dataclass
class JobContext:
    """Context for job processing."""
    job: JobRecord
    db: ClusterDB
    config: ClusterConfig


class JobQueue:
    """Job queue poller and processor."""

    def __init__(
        self,
        config: ClusterConfig,
        db: ClusterDB | None = None,
        poll_interval: int | None = None,
    ) -> None:
        self.config = config
        self.db = db
        self.poll_interval = poll_interval or config.poll_interval_seconds
        self._running = False
        self._task: asyncio.Task | None = None
        self._handlers: dict[str, JobHandler] = {}

    def register_handler(self, kind: str, handler: JobHandler) -> None:
        """Register a handler for a job kind."""
        self._handlers[kind] = handler
        logger.info("Registered handler for job kind: %s", kind)

    async def start(self) -> None:
        """Start polling for jobs."""
        if self.db is None:
            self.db = await get_cluster_db(self.config)
        self._running = True
        self._task = asyncio.create_task(self._poll_loop())
        logger.info("Job queue started (poll interval: %ds)", self.poll_interval)

    async def stop(self) -> None:
        """Stop polling."""
        self._running = False
        if self._task:
            self._task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self._task
        logger.info("Job queue stopped")

    async def _poll_loop(self) -> None:
        """Main polling loop."""
        while self._running:
            try:
                job = await self.db.claim_job(
                    node_id=self.config.node_id,
                    allowed_kinds=list(self._handlers.keys()) or self.config.allowed_job_kinds,
                    organization_id=self.config.organization_id,
                )
                if job:
                    await self._process_job(job)
                else:
                    await asyncio.sleep(self.poll_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error("Poll loop error: %s", e)
                await asyncio.sleep(self.poll_interval)

    async def _process_job(self, job: JobRecord) -> None:
        """Process a single job."""
        assert self.db is not None
        handler = self._handlers.get(job.kind)
        if not handler:
            logger.warning("No handler for job kind %s, requeuing", job.kind)
            await self.db.requeue_job(
                job.id, self.config.node_id, f"No handler for kind {job.kind}"
            )
            return

        logger.info("Processing job %s (kind: %s)", job.id, job.kind)
        ctx = JobContext(job=job, db=self.db, config=self.config)

        try:
            result = await handler(ctx)
            await self.db.complete_job(
                job_id=job.id,
                node_id=self.config.node_id,
                status="complete",
                result=result,
                artifacts=result.get("artifacts") if isinstance(result, dict) else None,
            )
            logger.info("Job %s completed", job.id)
        except Exception as e:
            logger.exception("Job %s failed: %s", job.id, e)
            await self.db.complete_job(
                job_id=job.id,
                node_id=self.config.node_id,
                status="failed",
                error=str(e),
            )


class JobHandler:
    """Base class for job handlers."""

    async def __call__(self, ctx: JobContext) -> dict[str, Any]:
        raise NotImplementedError


# Convenience functions for direct use

async def claim_job(
    config: ClusterConfig,
    db: ClusterDB,
    allowed_kinds: list[str] | None = None,
) -> JobRecord | None:
    """Claim a single job."""
    return await db.claim_job(
        node_id=config.node_id,
        allowed_kinds=allowed_kinds or config.allowed_job_kinds,
        organization_id=config.organization_id,
    )


async def complete_job(
    config: ClusterConfig,
    db: ClusterDB,
    job_id: uuid.UUID,
    status: str,
    result: dict[str, Any] | None = None,
    error: str | None = None,
    artifacts: list[dict[str, Any]] | None = None,
) -> bool:
    """Complete a job."""
    return await db.complete_job(
        job_id=job_id,
        node_id=config.node_id,
        status=status,
        result=result,
        error=error,
        artifacts=artifacts,
    )