"""OpenAI-compatible speech endpoint backed by ResembleAI Chatterbox."""

import gc
import os
import threading
import time
from io import BytesIO
from pathlib import Path

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


VOICES_DIR = Path(__file__).parent / "voices"
IDLE_UNLOAD_SECONDS = 300


class SpeechRequest(BaseModel):
    model: str = "chatterbox"
    input: str = Field(max_length=2000)
    voice: str
    response_format: str = "wav"
    speed: float = 1.0


class ModelManager:
    """Owns the GPU model and releases it after a period of inactivity."""

    def __init__(self) -> None:
        self._model = None
        self._unload_timer: threading.Timer | None = None
        self._lock = threading.Lock()
        self.last_request_time: float | None = None

    def _cancel_unload_timer(self) -> None:
        if self._unload_timer is not None:
            self._unload_timer.cancel()
            self._unload_timer = None

    def _schedule_unload(self) -> None:
        self._cancel_unload_timer()
        self._unload_timer = threading.Timer(IDLE_UNLOAD_SECONDS, self.unload)
        self._unload_timer.daemon = True
        self._unload_timer.start()

    def load_model(self):
        if self._model is None:
            from chatterbox.tts import ChatterboxTTS

            self._model = ChatterboxTTS.from_pretrained(device="cuda")
        return self._model

    def generate(self, text: str, voice_path: Path, speed: float) -> bytes:
        with self._lock:
            self._cancel_unload_timer()
            try:
                model = self.load_model()
                # Chatterbox does not expose a native speed control.  Keep the
                # OpenAI-compatible field accepted without changing speech speed.
                del speed
                audio = model.generate(text, audio_prompt_path=str(voice_path))

                import torchaudio

                buffer = BytesIO()
                torchaudio.save(buffer, audio, model.sr, format="wav")
                return buffer.getvalue()
            finally:
                if self._model is not None:
                    self.last_request_time = time.monotonic()
                    self._schedule_unload()

    def unload(self) -> None:
        with self._lock:
            self._cancel_unload_timer()
            if self._model is None:
                return
            self._model = None
            gc.collect()
            import torch

            torch.cuda.empty_cache()


def resolve_voice(voice: str) -> Path:
    candidate = Path(voice)
    if candidate.is_absolute():
        raise HTTPException(status_code=400, detail="voice must be a relative path")

    voices_root = VOICES_DIR.resolve()
    voice_path = Path(VOICES_DIR, voice).resolve()
    if not voice_path.is_relative_to(voices_root):
        raise HTTPException(status_code=400, detail="voice path escapes the voices directory")
    if not voice_path.is_file() or not os.access(voice_path, os.R_OK):
        raise HTTPException(status_code=400, detail="voice must be a readable regular file")
    return voice_path


app = FastAPI()
model_manager = ModelManager()


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/v1/models")
def list_models() -> dict[str, list[dict[str, str]]]:
    return {"data": [{"id": "chatterbox", "object": "model"}]}


@app.post("/v1/audio/speech")
def speech(request: SpeechRequest) -> Response:
    if request.response_format != "wav":
        raise HTTPException(status_code=400, detail="only wav response_format is supported")

    voice_path = resolve_voice(request.voice)
    audio = model_manager.generate(request.input, voice_path, request.speed)
    return Response(content=audio, media_type="audio/wav")