#!/usr/bin/env python3
"""Robo Pest STT relay — faster-whisper on nightmare's RTX 4080 SUPER.

POST /transcribe  (body: raw PCM 16kHz s16le mono bytes)  -> {"text": "..."}
GET  /health -> {"status": "ok", ...}

Auth: X-Auth-Token header required (set via STT_AUTH_TOKEN env var).
Health endpoint is open (no secrets).
"""
import io
import os
import struct
import numpy as np
from fastapi import FastAPI, Request, HTTPException
from faster_whisper import WhisperModel

MODEL_NAME = os.environ.get("STT_MODEL", "small.en")
DEVICE = os.environ.get("STT_DEVICE", "cuda")
AUTH_TOKEN = os.environ.get("STT_AUTH_TOKEN", "")

app = FastAPI()
model = WhisperModel(MODEL_NAME, device=DEVICE, compute_type="float16")


def check_auth(request: Request):
    if not AUTH_TOKEN:
        return
    token = request.headers.get("X-Auth-Token", "")
    if token != AUTH_TOKEN:
        raise HTTPException(status_code=401, detail="bad token")


def pcm16_to_float32(data: bytes) -> np.ndarray:
    n = len(data) // 2
    ints = struct.unpack(f"<{n}h", data[: n * 2])
    return np.array(ints, dtype=np.float32) / 32768.0


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


@app.post("/transcribe")
async def transcribe(request: Request):
    check_auth(request)
    body = await request.body()
    if not body:
        return {"text": ""}
    audio = pcm16_to_float32(body)
    segments, info = model.transcribe(audio, language="en", beam_size=1, vad_filter=True)
    text = " ".join(s.text.strip() for s in segments).strip()
    return {"text": text, "duration": round(info.duration, 2)}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8765)