#!/usr/bin/env python3
"""Small OpenAI-compatible Whisper service for the dedicated Jetson."""

from __future__ import annotations

import cgi
import json
import os
import tempfile
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import torch
import whisper


HOST = os.getenv("HERMES_STT_HOST", "0.0.0.0")
PORT = int(os.getenv("HERMES_STT_PORT", "9000"))
MODEL_NAME = os.getenv("HERMES_STT_MODEL", "large-v3-turbo")
CACHE_DIR = Path(os.getenv("HERMES_STT_CACHE", "/cache/whisper"))
MAX_AUDIO_BYTES = 30 * 1024 * 1024
MODEL_LOCK = threading.Lock()


def _json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.send_header("Cache-Control", "no-store")
    handler.end_headers()
    handler.wfile.write(body)


class SpeechHandler(BaseHTTPRequestHandler):
    """Serve health and transcription without exposing a general runtime."""

    server_version = "AtlasWhisper/1"

    def log_message(self, message: str, *args: object) -> None:
        print(f"[stt] {self.address_string()} {message % args}", flush=True)

    def do_GET(self) -> None:
        if self.path != "/health":
            _json(self, 404, {"error": "not found"})
            return
        _json(
            self,
            200,
            {
                "ok": True,
                "model": MODEL_NAME,
                "device": "cuda" if torch.cuda.is_available() else "cpu",
            },
        )

    def do_POST(self) -> None:
        if self.path != "/v1/audio/transcriptions":
            _json(self, 404, {"error": "not found"})
            return
        content_length = int(self.headers.get("Content-Length", "0") or "0")
        if content_length <= 0 or content_length > MAX_AUDIO_BYTES:
            _json(self, 413, {"error": "audio payload is missing or too large"})
            return

        content_type = self.headers.get("Content-Type", "")
        if not content_type.lower().startswith("multipart/form-data"):
            _json(self, 400, {"error": "multipart/form-data is required"})
            return

        form = cgi.FieldStorage(
            fp=self.rfile,
            headers=self.headers,
            environ={
                "REQUEST_METHOD": "POST",
                "CONTENT_TYPE": content_type,
                "CONTENT_LENGTH": str(content_length),
            },
        )
        audio = form["file"] if "file" in form else None
        if audio is None or not getattr(audio, "file", None):
            _json(self, 400, {"error": "file is required"})
            return

        suffix = Path(getattr(audio, "filename", "audio.wav") or "audio.wav").suffix
        suffix = suffix if suffix in {".wav", ".webm", ".ogg", ".mp3", ".m4a"} else ".wav"
        language = str(form.getfirst("language", "auto") or "auto").strip().lower()
        temp_path = ""
        try:
            with tempfile.NamedTemporaryFile(prefix="atlas-stt-", suffix=suffix, delete=False) as temp:
                temp_path = temp.name
                while True:
                    chunk = audio.file.read(1024 * 1024)
                    if not chunk:
                        break
                    temp.write(chunk)

            with MODEL_LOCK:
                result = self.server.model.transcribe(  # type: ignore[attr-defined]
                    temp_path,
                    language=None if language in {"", "auto"} else language,
                    task="transcribe",
                    fp16=torch.cuda.is_available(),
                    condition_on_previous_text=False,
                    temperature=0,
                    verbose=False,
                )
            transcript = str(result.get("text") or "").strip()
            _json(self, 200, {"text": transcript, "model": MODEL_NAME})
        except Exception as exc:
            print(f"[stt] transcription failed: {exc}", flush=True)
            _json(self, 500, {"error": "transcription failed"})
        finally:
            if temp_path:
                try:
                    os.unlink(temp_path)
                except OSError:
                    pass


def main() -> None:
    """Warm Whisper once, then serve concurrent clients through one GPU lock."""
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for the Atlas Whisper service")
    print(f"[stt] loading Whisper {MODEL_NAME} into CUDA", flush=True)
    model = whisper.load_model(MODEL_NAME, device="cuda", download_root=str(CACHE_DIR))
    server = ThreadingHTTPServer((HOST, PORT), SpeechHandler)
    server.model = model  # type: ignore[attr-defined]
    print(f"[stt] ready on {HOST}:{PORT}", flush=True)
    server.serve_forever(poll_interval=0.25)


if __name__ == "__main__":
    main()