src/midnight_static/asset_library.py

"""Load the published SFX/music asset library and resolve clips for a broadcast.

This wires the agentic retrieve-vs-generate matcher (`sfx.SFXLibrary` /
`music.MusicLibrary`) into the live pipeline: a real broadcast resolves its
scene SFX prompts and music plan against the published asset library, loads the
matched clips, and emits the same `SFXTrace`/`MusicTrace` shape as
`modal/traces.py` — so live broadcasts produce Best-Agent evidence.

Runtime matching uses the dependency-free `HashingEmbedder` (the manifest's
stored bge-small embeddings are dropped on load), so the live Space needs no
extra model. Music is NEVER generated at request time; if the library is
unavailable the mixer falls back to its synthetic bed/SFX (the station
improvises). No model API is called here — only static asset files are read.
"""

from __future__ import annotations

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

import numpy as np

from .mixer import SAMPLE_RATE, read_mono_wav
from .music import MusicAsset, MusicLibrary
from .schema import Genre, Script
from .sfx import SFXAsset, SFXLibrary

ASSETS_REPO = "build-small-hackathon/midnight-static-assets"
# Repo-local asset dir (populated by modal/batch_audio.py or a snapshot download).
LOCAL_ASSETS_DIR = Path(__file__).resolve().parent.parent.parent / "assets"


@dataclass
class AssetBundle:
    sfx: SFXLibrary
    music: MusicLibrary
    root: Path


@dataclass
class ResolvedAssets:
    sfx_layers: list[tuple[np.ndarray, float]] | None = None
    bed: np.ndarray | None = None
    opening_sting: np.ndarray | None = None
    closing_sting: np.ndarray | None = None
    traces: list[dict[str, Any]] = field(default_factory=list)
    source: str = "synthetic"  # "library" once real clips are loaded


def _manifest_root(assets_dir: Path | None) -> Path | None:
    """Find a directory containing manifest.json: explicit dir, repo assets/, or Hub."""
    for candidate in (assets_dir, LOCAL_ASSETS_DIR):
        if candidate and (candidate / "manifest.json").exists():
            return candidate
    try:  # last resort: pull the published asset dataset (static files, not a model)
        from huggingface_hub import snapshot_download

        path = snapshot_download(ASSETS_REPO, repo_type="dataset")
        if (Path(path) / "manifest.json").exists():
            return Path(path)
    except Exception:  # noqa: BLE001 - offline / no hub access → synthetic fallback
        return None
    return None


def load_asset_libraries(assets_dir: Path | None = None) -> AssetBundle | None:
    """Build SFX/music libraries from the manifest, or None if unavailable.

    Stored bge-small embeddings are dropped so the runtime `HashingEmbedder`
    re-embeds prompts (no torch/sentence-transformers needed at request time).
    """
    root = _manifest_root(assets_dir)
    if root is None:
        return None
    try:
        data = json.loads((root / "manifest.json").read_text(encoding="utf-8"))
    except Exception:  # noqa: BLE001
        return None
    entries = data.get("assets", data) if isinstance(data, dict) else data
    sfx_assets = [
        SFXAsset(id=str(e["id"]), prompt=str(e.get("prompt", "")), file=str(e.get("file", "")))
        for e in entries
        if str(e.get("kind", "")) == "sfx"
    ]
    music_assets = [
        MusicAsset(
            id=str(e["id"]),
            kind=str(e["kind"]),
            genre=Genre(e["genre"]),
            prompt=str(e.get("prompt", "")),
            file=str(e.get("file", "")),
        )
        for e in entries
        if str(e.get("kind", "")) in {"bed", "sting"}
    ]
    if not sfx_assets and not music_assets:
        return None
    return AssetBundle(SFXLibrary(sfx_assets), MusicLibrary(music_assets), root)


def load_clip(bundle: AssetBundle, asset) -> np.ndarray | None:
    """Read a matched asset's 24kHz mono WAV, or None if the file is missing."""
    if asset is None or not getattr(asset, "file", ""):
        return None
    path = bundle.root / asset.file
    if not path.exists():
        return None
    try:
        return read_mono_wav(path)
    except Exception:  # noqa: BLE001 - bad/missing wav → skip this layer
        return None


def resolve_broadcast_assets(
    script: Script,
    dialogue_seconds: float,
    bundle: AssetBundle | None = None,
) -> ResolvedAssets:
    """Resolve a script's SFX + music against the library, loading audio + traces."""
    if bundle is None:
        bundle = load_asset_libraries()
    if bundle is None:
        return ResolvedAssets(traces=[{"kind": "note", "action": "library_unavailable"}])

    traces: list[dict[str, Any]] = []
    sfx_layers: list[tuple[np.ndarray, float]] = []
    total_lines = sum(len(scene.lines) for scene in script.scenes) or 1
    lead_seconds = 0.8
    cumulative = 0
    for scene in script.scenes:
        # Place a scene's SFX just before its first line (proportional estimate).
        offset = lead_seconds + (cumulative / total_lines) * dialogue_seconds
        for prompt in scene.sfx:
            asset, trace = bundle.sfx.resolve(prompt)
            traces.append(
                {
                    "kind": "sfx",
                    "genre": script.genre.value,
                    "script": script.title,
                    "query": trace.query,
                    "action": trace.status,
                    "asset_id": trace.asset_id,
                    "score": trace.score,
                }
            )
            audio = load_clip(bundle, asset)
            if audio is not None:
                sfx_layers.append((audio, max(0.0, offset - 0.4)))
        cumulative += len(scene.lines)

    selection = bundle.music.select(script.music)
    for mt in selection.traces:
        traces.append(
            {
                "kind": "music",
                "genre": script.genre.value,
                "script": script.title,
                "role": mt.role,
                "query": mt.query,
                "action": "genre_fallback" if mt.genre_fallback else "in_genre_match",
                "asset_id": mt.asset_id,
                "score": mt.score,
            }
        )

    return ResolvedAssets(
        sfx_layers=sfx_layers or None,
        bed=load_clip(bundle, selection.bed),
        opening_sting=load_clip(bundle, selection.opening_sting),
        closing_sting=load_clip(bundle, selection.closing_sting),
        traces=traces,
        source="library",
    )


def write_traces(traces: list[dict[str, Any]], path: Path) -> Path:
    """Append-friendly JSONL trace sink (Best-Agent evidence from live runs)."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for trace in traces:
            handle.write(json.dumps(trace, ensure_ascii=False) + "\n")
    return path


def dialogue_duration_seconds(dialogue_wav: Path) -> float:
    try:
        return len(read_mono_wav(dialogue_wav)) / SAMPLE_RATE
    except Exception:  # noqa: BLE001
        return 60.0