"""HTTP API behind the web UI.

Encodes text queries with CLAP. The 2D projection and evaluation tables are
computed once and cached to disk.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from . import onnx_encoder
from .evaluate import embedding_health
from .voices import (
    FILTERABLE,
    VoiceCatalog,
    attribute_retrieval,
    facets,
    linear_probe,
    matching_indices,
    search_voices,
)

# Each query targets a label the catalog carries, so precision@10 is measurable.
PROBE_QUERIES = [
    ("gender", "female", "a recording of a woman speaking"),
    ("gender", "male", "a recording of a man speaking"),
    ("age", "young", "a recording of a young person in their twenties speaking"),
    ("age", "old", "a recording of an elderly person speaking"),
]

# About sound -- timbre, pace, energy -- because that is what CLAP encodes.
# Accent and language go through the filters; prompting for them is near random.
SUGGESTIONS = [
    "a warm raspy middle aged male narrator for an audiobook",
    "a deep authoritative movie trailer voice",
    "a gentle soothing voice for a meditation app",
    "a bright energetic young woman for a commercial",
    "a calm unhurried storyteller",
    "an intense dramatic villain",
]


class SearchRequest(BaseModel):
    query: str = Field(min_length=1, max_length=500)
    topN: int = Field(default=12, ge=1, le=50)
    # Applied before ranking. The library labels these exactly, and they are
    # hard constraints -- CLAP cannot hear accent or language at all, and hears
    # gender at 0.90-1.00, which is not a guarantee.
    filters: dict[str, str] = Field(default_factory=dict)


def _encoder(model_id: str, models: str | Path | None):
    """Prefer the ONNX export; fall back to torch.

    Serving only ever encodes text, so the int8 text tower is enough. It is
    lossy -- see mss.onnx_encoder -- and absent unless `mss export` has run.
    """
    if models and onnx_encoder.available(models):
        try:
            return onnx_encoder.OnnxTextEncoder(Path(models)), "onnx-int8"
        except ImportError:
            pass
    # Imported here, not at module scope: mss.clap pulls in torch, which is
    # ~200MB resident and the whole reason the ONNX path exists.
    from .clap import ClapEncoder

    return ClapEncoder(model_id=model_id), "torch"


class State:
    """Loaded catalog plus the derived data the UI needs."""

    def __init__(self, root: Path, models: str | Path | None = None) -> None:
        self.root = root
        self.catalog = VoiceCatalog.load(root)
        if not len(self.catalog) or self.catalog.vectors is None:
            raise RuntimeError(f"no voices at {root} -- run `mss fetch` first")
        self.encoder, self.backend = _encoder(self.catalog.model_id, models)
        self.by_id = {v.voice_id: i for i, v in enumerate(self.catalog.voices)}

    # -- cached derivations -------------------------------------------
    def _cached(self, name: str, build):
        path = self.root / f"{name}.json"
        if path.exists():
            return json.loads(path.read_text())
        data = build()
        path.write_text(json.dumps(data))
        return data

    def projection(self) -> dict:
        """2D t-SNE of the voice vectors, for the embedding-space view."""

        def build():
            from sklearn.manifold import TSNE

            xy = TSNE(n_components=2, perplexity=30, init="pca",
                      random_state=0).fit_transform(self.catalog.vectors)
            # normalise to [0, 1] so the UI can scale to any viewport
            lo, hi = xy.min(axis=0), xy.max(axis=0)
            xy = (xy - lo) / np.maximum(hi - lo, 1e-9)
            return {
                "points": [
                    {"id": v.voice_id, "name": v.name,
                     "x": round(float(p[0]), 4), "y": round(float(p[1]), 4),
                     "gender": v.gender, "age": v.age, "use_case": v.use_case}
                    for v, p in zip(self.catalog.voices, xy)
                ]
            }

        return self._cached("projection", build)

    def evaluation(self) -> dict:
        """Zero-shot, linear probe and retrieval precision, computed once."""

        def build():
            return {
                "probe": [linear_probe(self.catalog, a)
                          for a in ("gender", "age", "use_case", "language")],
                "retrieval": [attribute_retrieval(self.catalog, self.encoder, a, v, q)
                              for a, v, q in PROBE_QUERIES],
            }

        return self._cached("evaluation", build)


def create_app(root: str | Path, models: str | Path | None = None) -> FastAPI:
    state = State(Path(root).expanduser(), models)
    app = FastAPI(title="Voice Library Search", version="1.0")

    # The Vite dev server runs on a different origin.
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["http://localhost:5173", "http://127.0.0.1:5173"],
        allow_methods=["GET", "POST"],
        allow_headers=["*"],
    )

    def voice_payload(index: int, score: float | None = None, rank: int | None = None) -> dict:
        v = state.catalog.voices[index]
        return {
            "id": v.voice_id, "name": v.name, "gender": v.gender, "age": v.age,
            "accent": v.accent, "useCase": v.use_case, "language": v.language,
            "description": v.description,
            "audio": f"/api/voices/{v.voice_id}/audio",
            **({"score": round(score, 6)} if score is not None else {}),
            **({"rank": rank} if rank is not None else {}),
        }

    @app.get("/api/catalog")
    def catalog():
        from collections import Counter

        cat = state.catalog
        health = embedding_health(cat.vectors)
        return {
            "count": len(cat),
            "model": cat.model_id,
            "backend": state.backend,
            "dim": int(cat.vectors.shape[1]),
            "health": health,
            "suggestions": SUGGESTIONS,
            "filterable": list(FILTERABLE),
            "facets": {k: v for k, v in facets(cat).items()},
            "distributions": {
                attr: Counter(x for x in cat.labels(attr) if x).most_common(6)
                for attr in ("gender", "age", "language", "use_case")
            },
        }

    @app.post("/api/search")
    def search(req: SearchRequest):
        query = req.query.strip()
        if not query:
            raise HTTPException(400, "empty query")
        unknown = set(req.filters) - set(FILTERABLE)
        if unknown:
            raise HTTPException(400, f"unknown filter(s): {', '.join(sorted(unknown))}")

        # computed once and passed down so the catalog is not walked twice
        keep = matching_indices(state.catalog, req.filters)
        hits = search_voices(state.catalog, state.encoder, query,
                             top_n=req.topN, keep=keep)
        candidates = len(keep)
        results = [
            voice_payload(state.by_id[h.metadata["voice_id"]], h.score, h.rank)
            for h in hits
        ]
        return {
            "query": query,
            "total": len(state.catalog),
            "candidates": candidates,
            "filters": {k: v for k, v in req.filters.items() if v},
            "results": results,
        }

    @app.post("/api/heatmap")
    def heatmap(req: SearchRequest):
        """Similarity of every voice to the query, for shading the projection.

        Unlike /api/search this keeps all scores rather than the top N.
        """
        query = req.query.strip()
        if not query:
            raise HTTPException(400, "empty query")
        scores = state.catalog.vectors @ state.encoder.encode_text([query])[0]
        order = np.argsort(-scores)[: req.topN]
        return {
            "query": query,
            # aligned index-for-index with /api/projection points
            "scores": [round(float(s), 4) for s in scores],
            "min": round(float(scores.min()), 4),
            "max": round(float(scores.max()), 4),
            "top": [voice_payload(int(j), float(scores[j]), r)
                    for r, j in enumerate(order, 1)],
        }

    @app.get("/api/voices/{voice_id}/similar")
    def similar(voice_id: str, k: int = 6):
        if voice_id not in state.by_id:
            raise HTTPException(404, "unknown voice")
        i = state.by_id[voice_id]
        scores = state.catalog.vectors @ state.catalog.vectors[i]
        order = np.argsort(-scores)[1 : k + 1]        # [0] is the voice itself
        return {
            "seed": voice_payload(i),
            "results": [voice_payload(int(j), float(scores[j]), r)
                        for r, j in enumerate(order, 1)],
        }

    @app.get("/api/voices/{voice_id}/audio")
    def audio(voice_id: str):
        if voice_id not in state.by_id:
            raise HTTPException(404, "unknown voice")
        v = state.catalog.voices[state.by_id[voice_id]]
        path = state.catalog.preview_path(v)
        if path is None:
            raise HTTPException(404, "preview missing")
        return FileResponse(path, media_type="audio/mpeg")

    @app.get("/api/projection")
    def projection():
        return state.projection()

    @app.get("/api/evaluation")
    def evaluation():
        return state.evaluation()

    @app.get("/api/robustness")
    def robustness():
        """Benchmark results for the Evaluation tab, read from disk.

        Missing files return null so the UI can name the command to run.
        """
        out = {}
        for key, filename in (("compound", "compound_degradation.json"),):
            path = state.root / filename
            out[key] = json.loads(path.read_text()) if path.exists() else None
        return out

    app.state.backend = state.backend

    # Serve the built UI when it exists, so one process runs everything.
    dist = Path(__file__).resolve().parent.parent.parent / "ui" / "dist"
    if dist.is_dir():
        app.mount("/", StaticFiles(directory=dist, html=True), name="ui")

    return app


def serve(root: str | Path, host: str = "127.0.0.1", port: int = 8787,
          models: str | Path | None = None) -> None:
    import uvicorn

    app = create_app(root, models)
    dist = Path(__file__).resolve().parent.parent.parent / "ui" / "dist"
    print(f"Voice Library Search API  ->  http://{host}:{port}")
    print(f"text encoder: {app.state.backend}")
    print("UI: " + (f"http://{host}:{port}" if dist.is_dir()
                    else "run `npm run dev` in ui/ (http://localhost:5173)"))
    uvicorn.run(app, host=host, port=port, log_level="warning")