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

s2.cpp is a pure C++/GGML inference engine for Fish s2-pro GGUFs. It's
CLI-only (alpha), so this shim turns each HTTP POST into a subprocess
invocation that:

  1. Writes the optional reference WAV to /tmp/<id>.wav
  2. Runs:  s2 -m <model.gguf> -t <tokenizer.json>
              [-pa /tmp/ref.wav -pt "<transcript>"]
              -text "..." -o /tmp/<id>.out.wav
  3. Streams the resulting WAV back to the client.

The contract matches Fish's wrapper close enough that bench harnesses
written for fish-s2 work against fish-cpp without changes — same
POST /v1/tts, same {text, references[{audio,text}], format} request body.

What's NOT supported (vs Fish HF):
  * `streaming: true` — s2.cpp writes a complete WAV before returning.
    The shim accepts the field but ignores it. TTFB ≈ total wall time.
    If/when s2.cpp grows incremental output, wire it through here.
  * Paralinguistic tags — s2.cpp uses the same model weights, so tags
    SHOULD work, but quality may differ from the HF wrapper.
"""

from __future__ import annotations

import base64
import os
import subprocess
import tempfile
import uuid
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

S2_BINARY = os.environ.get("S2_BINARY", "/usr/local/bin/s2")
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")
DEVICE_INDEX = os.environ.get("FISH_CPP_DEVICE", "0")  # CUDA device id

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 (matches Fish HF schema)
    text: str   # transcript of the reference clip


class TTSRequest(BaseModel):
    text: str
    references: list[ReferenceAudio] = []
    format: str = "wav"
    streaming: bool = False  # accepted for contract 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}")
    return {"status": "ok", "model": MODEL_FILE}


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

    job = uuid.uuid4().hex
    out_path = Path(tempfile.gettempdir()) / f"fish-cpp-{job}.wav"
    ref_path: Optional[Path] = None

    cmd: list[str] = [
        S2_BINARY,
        "-m", str(MODEL_PATH),
        "-t", str(TOKENIZER_PATH),
        "-text", req.text,
        "-o", str(out_path),
        # -c <device> selects CUDA backend on the given device id. The
        # README example uses `-v 0` (which is --vulkan 0 — easy to
        # misread as "voice 0"); we want CUDA so that the work hits
        # the A6000 tensor cores instead of falling back to CPU.
        # Without a backend flag, s2 prints "NPU not compiled, falling
        # back to CPU" and runs at single-digit RTF on a 4-core CPU.
        "-c", "0",
    ]

    if req.references:
        # s2.cpp accepts a single reference at a time via -pa/-pt. If the
        # client sent multiple, use the first; the upstream API doesn't
        # appear to support multi-reference cloning.
        ref = req.references[0]
        ref_path = Path(tempfile.gettempdir()) / f"fish-cpp-{job}-ref.wav"
        ref_path.write_bytes(base64.b64decode(ref.audio))
        cmd.extend(["-pa", str(ref_path), "-pt", ref.text])

    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = DEVICE_INDEX

    try:
        result = subprocess.run(
            cmd, env=env, capture_output=True, text=True, timeout=180,
        )
    except subprocess.TimeoutExpired:
        raise HTTPException(504, "s2 binary timed out (180s)")
    finally:
        if ref_path and ref_path.exists():
            ref_path.unlink()

    if result.returncode != 0:
        raise HTTPException(
            500,
            f"s2 binary failed (rc={result.returncode}): {result.stderr[-500:]}",
        )

    if not out_path.exists() or out_path.stat().st_size == 0:
        raise HTTPException(500, "s2 binary produced no output")

    # FileResponse streams the file + cleans up after sending. Caller
    # gets WAV bytes immediately; we reap the temp file via background
    # task once the response finishes.
    from fastapi import BackgroundTasks
    bg = BackgroundTasks()
    bg.add_task(out_path.unlink, missing_ok=True)
    return FileResponse(
        out_path, media_type="audio/wav", filename=f"{job}.wav", background=bg,
    )