#!/usr/bin/env python3
# /// script
# requires-python = ">=3.9"
# dependencies = [
#   "faster-whisper>=1.0.0",
# ]
# ///
"""
NexusAI local STT sidecar (faster-whisper).

Invoked by the TypeScript MCP server as:
    uv run sidecar/stt_server.py '<json-request>'

Request JSON:
    { "audio": "/path/to/file", "model": "small", "device": "auto",
      "compute_type": "auto", "language": "ru" }

Prints exactly ONE JSON line on stdout:
    { "ok": true, "text": "...", "language": "ru", "duration": 12.3,
      "segments": [ { "start": 0.0, "end": 2.1, "text": "..." } ] }
On error:
    { "ok": false, "error": "message" }

All diagnostics go to stderr so stdout stays a clean single JSON line.
faster-whisper decodes audio via bundled ffmpeg/pyav, so no manual conversion
is required for common formats (mp3/wav/ogg/m4a/webm...).
"""
import json
import sys


def log(msg: str) -> None:
    print(msg, file=sys.stderr, flush=True)


def resolve_device(device: str, compute_type: str):
    if device and device != "auto":
        ct = compute_type if compute_type and compute_type != "auto" else (
            "float16" if device == "cuda" else "int8"
        )
        return device, ct
    # auto-detect: try cuda, fall back to cpu
    try:
        import ctranslate2  # noqa: F401
        cuda_count = 0
        try:
            cuda_count = ctranslate2.get_cuda_device_count()
        except Exception:
            cuda_count = 0
        if cuda_count > 0:
            return "cuda", (compute_type if compute_type != "auto" else "float16")
    except Exception:
        pass
    return "cpu", (compute_type if compute_type != "auto" else "int8")


def main() -> int:
    if len(sys.argv) < 2:
        print(json.dumps({"ok": False, "error": "missing request argument"}))
        return 1

    try:
        req = json.loads(sys.argv[1])
    except Exception as exc:
        print(json.dumps({"ok": False, "error": f"invalid request json: {exc}"}))
        return 1

    audio = req.get("audio")
    if not audio:
        print(json.dumps({"ok": False, "error": "no audio path provided"}))
        return 1

    model_size = req.get("model") or "small"
    device, compute_type = resolve_device(req.get("device", "auto"), req.get("compute_type", "auto"))
    language = req.get("language") or None

    try:
        from faster_whisper import WhisperModel
    except Exception as exc:
        print(json.dumps({"ok": False, "error": f"faster-whisper not available: {exc}"}))
        return 1

    # Allow pointing at a local model dir (offline) or a mirror via env.
    # NEXUS_STT_MODEL may be a whisper size name OR an absolute path to a
    # pre-downloaded CTranslate2 model directory.
    import os
    model_ref = model_size
    local_only = os.environ.get("NEXUS_STT_LOCAL_ONLY", "").lower() in ("1", "true", "yes")
    download_root = os.environ.get("NEXUS_STT_DOWNLOAD_ROOT") or None

    try:
        log(f"loading model={model_ref} device={device} compute={compute_type} local_only={local_only}")
        model = WhisperModel(
            model_ref,
            device=device,
            compute_type=compute_type,
            download_root=download_root,
            local_files_only=local_only,
        )
        segments_iter, info = model.transcribe(audio, language=language, vad_filter=True)

        segments = []
        parts = []
        for seg in segments_iter:
            parts.append(seg.text)
            segments.append({"start": round(seg.start, 3), "end": round(seg.end, 3), "text": seg.text.strip()})

        result = {
            "ok": True,
            "text": "".join(parts).strip(),
            "language": info.language,
            "duration": round(getattr(info, "duration", 0.0) or 0.0, 3),
            "segments": segments,
        }
        print(json.dumps(result, ensure_ascii=False))
        return 0
    except Exception as exc:
        hint = ""
        msg = str(exc)
        if "connect" in msg.lower() or "internet" in msg.lower() or "Hub" in msg:
            hint = (
                " | Model download failed. Options: set HF_ENDPOINT to a reachable mirror, "
                "pre-download the model and set NEXUS_STT_MODEL to its local path with "
                "NEXUS_STT_LOCAL_ONLY=1, or set NEXUS_STT_DOWNLOAD_ROOT to a cache dir."
            )
        print(json.dumps({"ok": False, "error": f"transcription failed: {exc}{hint}"}, ensure_ascii=False))
        return 1


if __name__ == "__main__":
    sys.exit(main())