UmbraAudio/scripts/cluster_worker.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""UmbraAudio cluster worker entry point.

Registers this node with UmbraTrans and polls for jobs.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import signal
from pathlib import Path
from typing import TYPE_CHECKING

import click

from umbraaudio.cluster import JobKind, JobQueue, load_cluster_config, register_node
from umbraaudio.cluster.database import close_cluster_db

if TYPE_CHECKING:
    from typing import Any

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)


@click.group()
def cli() -> None:
    """UmbraAudio cluster worker."""
    pass


@cli.command()
@click.option("--project-dir", type=click.Path(exists=True, path_type=Path), default=".")
@click.option("--once", is_flag=True, help="Claim and process one job then exit")
def worker(project_dir: Path, once: bool) -> None:
    """Run the cluster worker (register node, poll for jobs)."""
    asyncio.run(_run_worker(project_dir, once))


async def _run_worker(project_dir: Path, once: bool) -> None:
    config = load_cluster_config(project_dir)

    # Register node
    node = await register_node(config)

    # Create job queue
    queue = JobQueue(config, node.db)

    # Register job handlers
    from umbraaudio.cluster.jobs import JobContext, JobHandler

    class TTSHandler(JobHandler):
        async def __call__(self, ctx: JobContext) -> dict[str, Any]:
            logger.info("TTS job payload: %s", ctx.job.payload)
            # TODO: Implement actual TTS processing
            return {"status": "ok", "message": "TTS handler not yet implemented"}

    class AssembleHandler(JobHandler):
        async def __call__(self, ctx: JobContext) -> dict[str, Any]:
            logger.info("Assemble job payload: %s", ctx.job.payload)
            return {"status": "ok", "message": "Assemble handler not yet implemented"}

    queue.register_handler(JobKind.TTS, TTSHandler())
    queue.register_handler(JobKind.ASSEMBLE, AssembleHandler())

    # Handle shutdown signals
    shutdown = asyncio.Event()

    def signal_handler(sig, frame):
        logger.info("Received signal %s, shutting down...", sig)
        shutdown.set()

    for sig in (signal.SIGINT, signal.SIGTERM):
        with contextlib.suppress(OSError, ValueError):
            signal.signal(sig, signal_handler)

    try:
        if once:
            # Claim and process one job
            job = await queue.db.claim_job(
                node_id=config.node_id,
                allowed_kinds=list(queue._handlers.keys()),
                organization_id=config.organization_id,
            )
            if job:
                logger.info("Claimed job %s", job.id)
                # Process would happen here
            else:
                logger.info("No jobs available")
        else:
            # Run continuous polling
            await queue.start()
            await shutdown.wait()
            await queue.stop()
    finally:
        await node.stop()
        await close_cluster_db()


@cli.command()
@click.option("--project-dir", type=click.Path(exists=True, path_type=Path), default=".")
def check_cluster(project_dir: Path) -> None:
    """Check cluster connectivity (PostgreSQL, NFS)."""
    asyncio.run(_check_cluster(project_dir))


async def _check_cluster(project_dir: Path) -> None:
    config = load_cluster_config(project_dir)

    # Check PostgreSQL
    from umbraaudio.cluster.database import ClusterDB
    db = ClusterDB(config)
    try:
        await db.connect()
        async with db.acquire() as conn:
            version = await conn.fetchval("SELECT version()")
            logger.info("PostgreSQL OK: %s", version.split(",")[0])

            # Check tables
            tables = await conn.fetch(
                "SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename"
            )
            logger.info("Tables: %s", [t["tablename"] for t in tables])
    except Exception as e:
        logger.error("PostgreSQL check failed: %s", e)
    finally:
        await db.close()

    # Check NFS
    from umbraaudio.cluster.nfs import check_nfs_server, ensure_nfs_mounts
    if check_nfs_server(config):
        results = ensure_nfs_mounts(config)
        for name, ok in results.items():
            status = "OK" if ok else "FAILED"
            logger.info("NFS %s: %s", name, status)
    else:
        logger.error("NFS server check failed")


@cli.command()
@click.option("--project-dir", type=click.Path(exists=True, path_type=Path), default=".")
def register_node_cmd(project_dir: Path) -> None:
    """Register this node with the cluster (one-time)."""
    asyncio.run(_register_node_cmd(project_dir))


async def _register_node_cmd(project_dir: Path) -> None:
    config = load_cluster_config(project_dir)
    node = await register_node(config)
    logger.info("Node registered: %s", config.node_id)
    await node.stop()
    await close_cluster_db()


if __name__ == "__main__":
    cli()