UmbraAudio/umbraaudio/cluster/artifacts.py

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

"""Artifact protocol (content-addressed storage) for UmbraAudio ↔ UmbraTrans."""

from __future__ import annotations

import hashlib
import logging
import os
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, BinaryIO

if TYPE_CHECKING:
    from .config import ClusterConfig
    from .database import ClusterDB

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class ArtifactRef:
    """Reference to a written artifact."""
    artifact_id: uuid.UUID
    content_hash: str
    byte_size: int
    storage_uri: str
    local_path: Path


class ArtifactStore:
    """Content-addressed artifact storage on NFS."""

    def __init__(self, config: ClusterConfig, db: ClusterDB) -> None:
        self.config = config
        self.db = db

    def _artifact_path(self, org_id: uuid.UUID, project_id: uuid.UUID, content_hash: str) -> Path:
        """Compute artifact path on NFS."""
        # Structure: /mnt/umbra/artifacts/{org_id}/{project_id}/{hash[:2]}/{hash}
        artifacts_root = self.config.local_mount_root / "artifacts"
        return artifacts_root / str(org_id) / str(project_id) / content_hash[:2] / content_hash

    async def write(
        self,
        data: bytes,
        project_id: uuid.UUID,
        organization_id: uuid.UUID,
        lang: str,
        kind: str,
        media_type: str,
        derivation: dict[str, Any] | None = None,
    ) -> ArtifactRef:
        """Write artifact atomically to NFS and register in database."""
        # 1. Compute content hash
        content_hash = hashlib.sha256(data).hexdigest()
        byte_size = len(data)

        # 2. Write to NFS (atomic: tmp → rename)
        artifact_path = self._artifact_path(organization_id, project_id, content_hash)
        artifact_path.parent.mkdir(parents=True, exist_ok=True)

        tmp_path = artifact_path.with_suffix(artifact_path.suffix + ".tmp")
        try:
            with open(tmp_path, "wb") as f:
                f.write(data)
            os.rename(tmp_path, artifact_path)
        except Exception:
            if tmp_path.exists():
                tmp_path.unlink(missing_ok=True)
            raise

        # 3. Insert into database
        storage_key = content_hash
        storage_uri = f"file://{self.config.node_id}/{artifact_path.relative_to(self.config.local_mount_root)}"

        record = await self.db.insert_artifact(
            project_id=project_id,
            organization_id=organization_id,
            lang=lang,
            kind=kind,
            content_hash=content_hash,
            byte_size=byte_size,
            media_type=media_type,
            storage_key=storage_key,
            storage_uri=storage_uri,
            derivation=derivation,
        )

        logger.info(
            "Wrote artifact %s (%s, %d bytes) to %s",
            record.id, kind, byte_size, artifact_path
        )

        return ArtifactRef(
            artifact_id=record.id,
            content_hash=content_hash,
            byte_size=byte_size,
            storage_uri=storage_uri,
            local_path=artifact_path,
        )

    async def write_stream(
        self,
        stream: BinaryIO,
        project_id: uuid.UUID,
        organization_id: uuid.UUID,
        lang: str,
        kind: str,
        media_type: str,
        derivation: dict[str, Any] | None = None,
        chunk_size: int = 8192,
    ) -> ArtifactRef:
        """Write artifact from a stream (for large files)."""
        # For streaming, we need to compute hash while writing
        hasher = hashlib.sha256()
        temp_path = self.config.local_mount_root / "artifacts" / ".tmp" / f"{uuid.uuid4()}.tmp"
        temp_path.parent.mkdir(parents=True, exist_ok=True)

        try:
            byte_size = 0
            with open(temp_path, "wb") as f:
                while chunk := stream.read(chunk_size):
                    hasher.update(chunk)
                    f.write(chunk)
                    byte_size += len(chunk)

            content_hash = hasher.hexdigest()
            final_path = self._artifact_path(organization_id, project_id, content_hash)
            final_path.parent.mkdir(parents=True, exist_ok=True)

            # Atomic move
            os.rename(temp_path, final_path)

            storage_key = content_hash
            storage_uri = f"file://{self.config.node_id}/{final_path.relative_to(self.config.local_mount_root)}"

            record = await self.db.insert_artifact(
                project_id=project_id,
                organization_id=organization_id,
                lang=lang,
                kind=kind,
                content_hash=content_hash,
                byte_size=byte_size,
                media_type=media_type,
                storage_key=storage_key,
                storage_uri=storage_uri,
                derivation=derivation,
            )

            return ArtifactRef(
                artifact_id=record.id,
                content_hash=content_hash,
                byte_size=byte_size,
                storage_uri=storage_uri,
                local_path=final_path,
            )
        finally:
            if temp_path.exists():
                temp_path.unlink(missing_ok=True)

    async def read(self, artifact_id: uuid.UUID) -> bytes | None:
        """Read artifact by ID."""
        record = await self.db.get_artifact(artifact_id)
        if not record:
            return None

        # Parse storage_uri to get local path
        # file://node-id/mnt/umbra/artifacts/...
        if not record.storage_uri.startswith("file://"):
            logger.error("Unsupported storage URI scheme: %s", record.storage_uri)
            return None

        # Extract path after hostname
        path_str = record.storage_uri.split("/", 3)[-1]  # "/mnt/umbra/artifacts/..."
        local_path = Path(path_str)

        if not local_path.exists():
            logger.error("Artifact file not found: %s", local_path)
            return None

        with open(local_path, "rb") as f:
            data = f.read()

        # Verify hash
        if hashlib.sha256(data).hexdigest() != record.content_hash:
            logger.error("Artifact hash mismatch: %s", artifact_id)
            return None

        return data

    async def read_stream(self, artifact_id: uuid.UUID, chunk_size: int = 8192):
        """Read artifact as stream."""
        record = await self.db.get_artifact(artifact_id)
        if not record:
            return

        if not record.storage_uri.startswith("file://"):
            return

        path_str = record.storage_uri.split("/", 3)[-1]
        local_path = Path(path_str)

        if not local_path.exists():
            return

        with open(local_path, "rb") as f:
            while chunk := f.read(chunk_size):
                yield chunk

    def local_path(self, artifact_id: uuid.UUID) -> Path | None:
        """Get local path for artifact (without reading)."""
        # This would need a DB lookup; for now return None
        # Could be cached
        return None


async def write_artifact(
    config: ClusterConfig,
    db: ClusterDB,
    data: bytes,
    project_id: uuid.UUID,
    organization_id: uuid.UUID,
    lang: str,
    kind: str,
    media_type: str,
    derivation: dict[str, Any] | None = None,
) -> ArtifactRef:
    """Write artifact (convenience function)."""
    store = ArtifactStore(config, db)
    return await store.write(data, project_id, organization_id, lang, kind, media_type, derivation)


async def read_artifact(
    config: ClusterConfig,
    db: ClusterDB,
    artifact_id: uuid.UUID,
) -> bytes | None:
    """Read artifact (convenience function)."""
    store = ArtifactStore(config, db)
    return await store.read(artifact_id)