"""fish-cpp — HTTP shim wrapping s2.cpp's built-in server in Fish's
/v1/tts API contract.

Architecture: an `s2 --server` instance runs IN THE SAME CONTAINER on
localhost:3030 with the model loaded once on the GPU; this FastAPI
shim translates each Fish-shape JSON POST into a multipart form POST
against localhost:3030/generate and streams the WAV back.

The earlier subprocess-per-request design paid CUDA init + model load
(~10-20s) on every call, which dominated wall time. Keeping s2 server
resident moves the cost to startup (paid once) so per-request latency
matches what the model can actually do.

The shim still accepts (and ignores) `streaming: true` in the request
body for contract parity with fish-s2 — s2.cpp's /generate is
synchronous, so TTFB ≈ total wall time. If/when s2.cpp grows
incremental output, wire it through here.
"""

from __future__ import annotations

import base64
import io
import os
import tempfile
from pathlib import Path
from typing import Optional

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel

S2_BASE_URL = os.environ.get("S2_BASE_URL", "http://127.0.0.1:3030")
WEIGHTS_DIR = Path(os.environ.get("WEIGHTS_DIR", "/weights"))
MODEL_FILE = os.environ.get("FISH_CPP_MODEL", "s2-pro-q6_k.gguf")
TOKENIZER_FILE = os.environ.get("FISH_CPP_TOKENIZER", "tokenizer.json")

MODEL_PATH = WEIGHTS_DIR / MODEL_FILE
TOKENIZER_PATH = WEIGHTS_DIR / TOKENIZER_FILE

app = FastAPI(title="fish-cpp")


class ReferenceAudio(BaseModel):
    audio: str  # base64-encoded WAV bytes
    text: str   # transcript of the reference clip


class TTSRequest(BaseModel):
    text: str
    references: list[ReferenceAudio] = []
    format: str = "wav"
    streaming: bool = False  # accepted for parity; ignored


@app.get("/v1/health")
def health() -> dict:
    if not MODEL_PATH.exists():
        raise HTTPException(503, f"model missing: {MODEL_PATH}")
    if not TOKENIZER_PATH.exists():
        raise HTTPException(503, f"tokenizer missing: {TOKENIZER_PATH}")
    # Probe the s2 server backend — readiness, not just liveness.
    try:
        r = httpx.get(f"{S2_BASE_URL}/", timeout=2.0)
    except httpx.HTTPError as exc:
        raise HTTPException(503, f"s2 server unreachable: {exc}")
    return {"status": "ok", "model": MODEL_FILE, "s2_status": r.status_code}


@app.post("/v1/tts")
def tts(req: TTSRequest) -> Response:
    if req.format not in ("wav",):
        raise HTTPException(400, f"only wav format supported, got {req.format}")

    # Translate Fish-shaped JSON to s2.cpp's multipart form.
    files: dict[str, tuple] = {
        "text": (None, req.text),
    }
    if req.references:
        # s2.cpp accepts a single reference at a time. If the client
        # sent multiple, use the first.
        ref = req.references[0]
        files["prompt_text"] = (None, ref.text)
        files["prompt_audio"] = ("ref.wav", base64.b64decode(ref.audio), "audio/wav")

    try:
        # Long timeout — s2 generation is sub-realtime on q6_k. The
        # /v1/health probe upstream already vetted s2 server is live.
        r = httpx.post(
            f"{S2_BASE_URL}/generate", files=files, timeout=180.0,
        )
    except httpx.HTTPError as exc:
        raise HTTPException(502, f"s2 server request failed: {exc}")

    if r.status_code != 200:
        # s2 returns JSON `{"error": "..."}` on failure
        raise HTTPException(
            500, f"s2 server returned {r.status_code}: {r.text[:500]}"
        )

    return Response(
        content=r.content,
        media_type=r.headers.get("content-type", "audio/wav"),
    )