UmbraAudio/umbraaudio/cluster/config.py

"""Cluster configuration for UmbraAudio ↔ UmbraTrans integration."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Literal

from pydantic_settings import BaseSettings, SettingsConfigDict


@dataclass(frozen=True)
class ClusterConfig:
    """Runtime configuration for cluster integration."""

    # PostgreSQL (PgBouncer pooled)
    pg_host: str = "192.168.123.14"
    pg_port: int = 6432
    pg_database: str = "umbratrans"
    pg_user: str = "umbraudio"
    pg_password: str = ""
    pg_sslmode: Literal[
        "disable", "allow", "prefer", "require", "verify-ca", "verify-full"
    ] = "require"
    pg_app_name: str = "umbraudio-worker"

    # Direct PostgreSQL (admin/backup)
    pg_direct_port: int = 5432

    # NFS
    nfs_server: str = "192.168.123.14"
    nfs_productions_path: str = "/mnt/umbra-data/productions"
    nfs_artifacts_path: str = "/mnt/umbra-data/artifacts"
    nfs_covers_path: str = "/mnt/umbra-data/covers"
    nfs_audio_path: str = "/mnt/umbra-data/audio"

    # Local mount points
    local_mount_root: Path = Path("/mnt/umbra")

    # Compute node
    node_id: str = "umbra-audio-01"
    organization_id: str = "umbra"
    node_name: str = ""
    datacenter: str = "local"
    agent_version: str = "1.0.0"

    # Credentials
    node_credential_prefix: str = "umbra-audio-01-"
    node_scopes: tuple[str, ...] = ("queue:claim", "queue:complete", "artifacts:write")

    # Job polling
    poll_interval_seconds: int = 10
    lease_duration_minutes: int = 30
    allowed_job_kinds: tuple[str, ...] = (
        "tts", "asr", "assemble", "encode", "master", "speaker_profile",
    )

    # Heartbeat
    heartbeat_interval_seconds: int = 30
    heartbeat_timeout_minutes: int = 2

    @property
    def pg_dsn(self) -> str:
        """PostgreSQL DSN for psycopg/asyncpg."""
        return (
            f"postgresql://{self.pg_user}:{self.pg_password}@{self.pg_host}:{self.pg_port}/"
            f"{self.pg_database}?sslmode={self.pg_sslmode}&application_name={self.pg_app_name}"
        )

    @property
    def pg_direct_dsn(self) -> str:
        """Direct PostgreSQL DSN (port 5432)."""
        return (
            f"postgresql://{self.pg_user}:{self.pg_password}@{self.pg_host}:{self.pg_direct_port}/"
            f"{self.pg_database}?sslmode={self.pg_sslmode}&application_name={self.pg_app_name}-direct"
        )

    @property
    def nfs_mounts(self) -> dict[str, tuple[str, Path]]:
        """NFS exports → local mount points."""
        server = self.nfs_server
        root = self.local_mount_root
        return {
            "productions": (
                f"{server}:{self.nfs_productions_path}",
                root / "productions",
            ),
            "artifacts": (
                f"{server}:{self.nfs_artifacts_path}",
                root / "artifacts",
            ),
            "covers": (
                f"{server}:{self.nfs_covers_path}",
                root / "covers",
            ),
            "audio": (
                f"{server}:{self.nfs_audio_path}",
                root / "audio",
            ),
        }


class ClusterSettings(BaseSettings):
    """Settings loaded from environment (.env)."""

    model_config = SettingsConfigDict(env_prefix="UMBRA_CLUSTER__", env_file=".env", extra="ignore")

    pg_host: str = "192.168.123.14"
    pg_port: int = 6432
    pg_database: str = "umbratrans"
    pg_user: str = "umbraudio"
    pg_password: str = ""
    pg_sslmode: str = "require"
    pg_direct_port: int = 5432

    nfs_server: str = "192.168.123.14"
    nfs_productions_path: str = "/mnt/umbra-data/productions"
    nfs_artifacts_path: str = "/mnt/umbra-data/artifacts"
    nfs_covers_path: str = "/mnt/umbra-data/covers"
    nfs_audio_path: str = "/mnt/umbra-data/audio"

    local_mount_root: Path = Path("/mnt/umbra")

    node_id: str = "umbra-audio-01"
    organization_id: str = "umbra"
    datacenter: str = "local"
    agent_version: str = "1.0.0"

    node_credential_prefix: str = "umbra-audio-01-"
    node_scopes: str = "queue:claim,queue:complete,artifacts:write"

    poll_interval_seconds: int = 10
    lease_duration_minutes: int = 30
    allowed_job_kinds: str = "tts,asr,assemble,encode,master,speaker_profile"

    heartbeat_interval_seconds: int = 30
    heartbeat_timeout_minutes: int = 2


def load_cluster_config(project_dir: Path | None = None) -> ClusterConfig:
    """Load cluster configuration from environment and defaults."""
    settings = ClusterSettings()

    # Parse comma-separated strings
    scopes = tuple(s.strip() for s in settings.node_scopes.split(",") if s.strip())
    job_kinds = tuple(k.strip() for k in settings.allowed_job_kinds.split(",") if k.strip())

    # Node name defaults to hostname if not set
    node_name = settings.node_id  # Could be overridden

    return ClusterConfig(
        pg_host=settings.pg_host,
        pg_port=settings.pg_port,
        pg_database=settings.pg_database,
        pg_user=settings.pg_user,
        pg_password=settings.pg_password,
        pg_sslmode=settings.pg_sslmode,  # type: ignore[assignment]
        pg_direct_port=settings.pg_direct_port,
        nfs_server=settings.nfs_server,
        nfs_productions_path=settings.nfs_productions_path,
        nfs_artifacts_path=settings.nfs_artifacts_path,
        nfs_covers_path=settings.nfs_covers_path,
        nfs_audio_path=settings.nfs_audio_path,
        local_mount_root=settings.local_mount_root,
        node_id=settings.node_id,
        organization_id=settings.organization_id,
        node_name=node_name,
        datacenter=settings.datacenter,
        agent_version=settings.agent_version,
        node_credential_prefix=settings.node_credential_prefix,
        node_scopes=scopes,
        poll_interval_seconds=settings.poll_interval_seconds,
        lease_duration_minutes=settings.lease_duration_minutes,
        allowed_job_kinds=job_kinds,
        heartbeat_interval_seconds=settings.heartbeat_interval_seconds,
        heartbeat_timeout_minutes=settings.heartbeat_timeout_minutes,
    )