"""Music bed/sting library interface with selection traces.

Per ARCHITECTURE.md, music is ALWAYS chosen from a prebuilt library and never
generated at request time (beds are too slow to synthesize live). Given a
script's `MusicPlan`, pick the nearest bed and opening/closing stings within
the script's genre, recording a `MusicTrace` for each pick.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from .embeddings import Embedder, HashingEmbedder, cosine_similarity
from .schema import Genre, MusicPlan


@dataclass(frozen=True)
class MusicAsset:
    id: str
    kind: str  # "bed" | "sting"
    genre: Genre
    prompt: str
    file: str = ""
    embedding: tuple[float, ...] | None = None


@dataclass(frozen=True)
class MusicTrace:
    role: str  # "opening_sting" | "bed" | "closing_sting"
    query: str
    asset_id: str
    score: float
    genre_fallback: bool = False


@dataclass(frozen=True)
class MusicSelection:
    genre: Genre
    opening_sting: MusicAsset | None
    bed: MusicAsset | None
    closing_sting: MusicAsset | None
    traces: list[MusicTrace]


class MusicLibrary:
    """Nearest-match music selector over a fixed bed/sting library."""

    def __init__(
        self,
        assets: list[MusicAsset] | None = None,
        embedder: Embedder | None = None,
    ) -> None:
        self.embedder = embedder or HashingEmbedder()
        self.assets: list[MusicAsset] = []
        self._vectors: list[np.ndarray] = []
        for asset in assets or []:
            self._register(asset)

    def _register(self, asset: MusicAsset) -> None:
        vector = (
            np.asarray(asset.embedding, dtype=np.float32)
            if asset.embedding is not None
            else self.embedder.embed(asset.prompt)
        )
        self.assets.append(asset)
        self._vectors.append(vector)

    def _nearest(
        self, role: str, query: str, kind: str, genre: Genre
    ) -> tuple[MusicAsset | None, MusicTrace | None]:
        candidates = [
            (asset, vec)
            for asset, vec in zip(self.assets, self._vectors)
            if asset.kind == kind
        ]
        in_genre = [(a, v) for a, v in candidates if a.genre == genre]
        pool = in_genre or candidates
        if not pool:
            return None, None
        query_vec = self.embedder.embed(query)
        scores = [cosine_similarity(query_vec, vec) for _, vec in pool]
        best = int(np.argmax(scores))
        asset = pool[best][0]
        trace = MusicTrace(
            role=role,
            query=query,
            asset_id=asset.id,
            score=round(float(scores[best]), 4),
            genre_fallback=not in_genre,
        )
        return asset, trace

    def select(self, plan: MusicPlan) -> MusicSelection:
        """Choose opening sting, bed, and closing sting for a script's plan."""
        traces: list[MusicTrace] = []
        picks: dict[str, MusicAsset | None] = {}
        for role, query, kind in (
            ("opening_sting", plan.opening_sting, "sting"),
            ("bed", plan.bed, "bed"),
            ("closing_sting", plan.closing_sting, "sting"),
        ):
            asset, trace = self._nearest(role, query, kind, plan.genre)
            picks[role] = asset
            if trace is not None:
                traces.append(trace)
        return MusicSelection(
            genre=plan.genre,
            opening_sting=picks["opening_sting"],
            bed=picks["bed"],
            closing_sting=picks["closing_sting"],
            traces=traces,
        )

    @classmethod
    def from_manifest(
        cls, path: str | Path, embedder: Embedder | None = None
    ) -> "MusicLibrary":
        data = json.loads(Path(path).read_text(encoding="utf-8"))
        entries = data.get("assets", data) if isinstance(data, dict) else data
        assets = [
            MusicAsset(
                id=str(entry["id"]),
                kind=str(entry["kind"]),
                genre=Genre(entry["genre"]),
                prompt=str(entry.get("prompt", "")),
                file=str(entry.get("file", "")),
                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", "")) in {"bed", "sting"}
        ]
        return cls(assets, embedder)