src/midnight_static/sfx.py

"""SFX retrieval/generation interface with resolution traces.

Implements the ARCHITECTURE.md policy: embed a scene's SFX prompt, cosine-match
against a manifest, use the cached clip when the score clears a threshold,
otherwise generate a new clip and append it so the library learns. Every
resolution returns a `SFXTrace` the UI/loader can surface or log.
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol

import numpy as np

from .embeddings import Embedder, HashingEmbedder, cosine_similarity

DEFAULT_MATCH_THRESHOLD = 0.62


@dataclass(frozen=True)
class SFXAsset:
    id: str
    prompt: str
    file: str = ""
    kind: str = "sfx"
    embedding: tuple[float, ...] | None = None


@dataclass(frozen=True)
class SFXTrace:
    """A record of how one SFX prompt was resolved."""

    query: str
    status: str  # "cache_hit" | "generated"
    asset_id: str
    score: float

    @property
    def is_hit(self) -> bool:
        return self.status == "cache_hit"


class SFXGenerator(Protocol):
    """Produces a new SFX asset for a cache miss (real path: Stable Audio Open)."""

    def generate(self, prompt: str) -> SFXAsset: ...


class StubSFXGenerator:
    """Dependency-free generator that registers a deferred clip for a prompt."""

    def __init__(self) -> None:
        self._count = 0

    def generate(self, prompt: str) -> SFXAsset:
        self._count += 1
        asset_id = f"gen-{self._count:03d}"
        return SFXAsset(id=asset_id, prompt=prompt, file="", kind="sfx")


class SFXLibrary:
    """Embedding-matched SFX cache that generates and learns on a miss."""

    def __init__(
        self,
        assets: list[SFXAsset] | None = None,
        embedder: Embedder | None = None,
        *,
        threshold: float = DEFAULT_MATCH_THRESHOLD,
        generator: SFXGenerator | None = None,
    ) -> None:
        self.embedder = embedder or HashingEmbedder()
        self.threshold = threshold
        self.generator = generator or StubSFXGenerator()
        self.assets: list[SFXAsset] = []
        self._vectors: list[np.ndarray] = []
        for asset in assets or []:
            self._register(asset)

    def _register(self, asset: SFXAsset) -> SFXAsset:
        if asset.embedding is not None:
            vector = np.asarray(asset.embedding, dtype=np.float32)
        else:
            vector = self.embedder.embed(asset.prompt)
            asset = SFXAsset(
                id=asset.id,
                prompt=asset.prompt,
                file=asset.file,
                kind=asset.kind,
                embedding=tuple(float(x) for x in vector),
            )
        self.assets.append(asset)
        self._vectors.append(vector)
        return asset

    def best_match(self, prompt: str) -> tuple[SFXAsset | None, float]:
        if not self.assets:
            return None, 0.0
        query = self.embedder.embed(prompt)
        scores = [cosine_similarity(query, vector) for vector in self._vectors]
        best_index = int(np.argmax(scores))
        return self.assets[best_index], float(scores[best_index])

    def resolve(self, prompt: str) -> tuple[SFXAsset, SFXTrace]:
        """Return the chosen asset and a trace; generate-and-learn on a miss."""
        match, score = self.best_match(prompt)
        if match is not None and score >= self.threshold:
            return match, SFXTrace(prompt, "cache_hit", match.id, round(score, 4))
        generated = self._register(self.generator.generate(prompt))
        return generated, SFXTrace(prompt, "generated", generated.id, round(score, 4))

    def resolve_all(self, prompts: list[str]) -> tuple[list[SFXAsset], list[SFXTrace]]:
        assets: list[SFXAsset] = []
        traces: list[SFXTrace] = []
        for prompt in prompts:
            asset, trace = self.resolve(prompt)
            assets.append(asset)
            traces.append(trace)
        return assets, traces

    @classmethod
    def from_manifest(
        cls,
        path: str | Path,
        embedder: Embedder | None = None,
        **kwargs,
    ) -> "SFXLibrary":
        data = json.loads(Path(path).read_text(encoding="utf-8"))
        entries = data.get("assets", data) if isinstance(data, dict) else data
        assets = [
            SFXAsset(
                id=str(entry["id"]),
                prompt=str(entry.get("prompt", "")),
                file=str(entry.get("file", "")),
                kind=str(entry.get("kind", "sfx")),
                embedding=(
                    tuple(float(x) for x in entry["embedding"])
                    if entry.get("embedding") is not None
                    else None
                ),
            )
            for entry in entries
            if str(entry.get("kind", "sfx")) == "sfx"
        ]
        return cls(assets, embedder, **kwargs)