#!/usr/bin/env python3
"""OpenAI-compatible STT wrapper around faster-whisper, for Open Notebook.

Implements exactly the endpoint esperanto's OpenAICompatibleSpeechToTextModel calls:
    POST /audio/transcriptions  (multipart: file, model, language?, prompt?) -> {"text": ...}

Start:
    python3 stt_server.py

Env vars:
    STT_HOST (default 0.0.0.0), STT_PORT (default 8902)
    WHISPER_MODEL (default large-v3), WHISPER_COMPUTE_TYPE (default float16)
    CUDA_VISIBLE_DEVICES should be set by the caller (e.g. "1,2") to keep GPU 0 free.
"""
from __future__ import annotations

import os
import tempfile
from pathlib import Path

from faster_whisper import WhisperModel
from fastapi import FastAPI, File, Form, UploadFile

app = FastAPI(title="faster-whisper STT (OpenAI-compatible)", version="1.0")

_MODEL_SIZE = os.environ.get("WHISPER_MODEL", "large-v3")
_COMPUTE_TYPE = os.environ.get("WHISPER_COMPUTE_TYPE", "float16")
_model = WhisperModel(_MODEL_SIZE, device="cuda", compute_type=_COMPUTE_TYPE)


@app.get("/health")
def health():
    return {"status": "ok", "model": _MODEL_SIZE}


@app.post("/audio/transcriptions")
async def transcriptions(
    file: UploadFile = File(...),
    model: str = Form(default="whisper-1"),
    language: str | None = Form(default=None),
    prompt: str | None = Form(default=None),
):
    suffix = Path(file.filename or "audio").suffix or ".wav"
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(await file.read())
        tmp_path = tmp.name

    try:
        segments, _info = _model.transcribe(
            tmp_path,
            language=language,
            initial_prompt=prompt,
        )
        text = "".join(segment.text for segment in segments).strip()
    finally:
        os.unlink(tmp_path)

    return {"text": text}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=os.environ.get("STT_HOST", "0.0.0.0"),
                port=int(os.environ.get("STT_PORT", "8902")))