#!/usr/bin/env python3
"""CUDA-only Voxtral ASR worker managed by Omnius."""

from __future__ import annotations

import argparse
import json
import os
import platform
import sys
import time
from pathlib import Path
from typing import Any

MODELS = {
    "voxtral-mini-4b-realtime-2602": {
        "upstream": "mistralai/Voxtral-Mini-4B-Realtime-2602",
        "realtime": True,
        "minimum_vram_gb": 15.0,
    },
    "voxtral-mini-3b-2507": {
        "upstream": "mistralai/Voxtral-Mini-3B-2507",
        "realtime": False,
        "minimum_vram_gb": 9.0,
    },
    "voxtral-small-24b-2507": {
        "upstream": "mistralai/Voxtral-Small-24B-2507",
        "realtime": False,
        "minimum_vram_gb": 54.0,
    },
}


def emit(event: dict[str, Any]) -> None:
    sys.stdout.write(json.dumps(event, ensure_ascii=True) + "\n")
    sys.stdout.flush()


def error(message: str) -> None:
    emit({"type": "error", "message": message})


def model_path(model_id: str) -> Path:
    root = Path(
        os.environ.get(
            "OMNIUS_ASR_MODEL_DIR",
            str(Path.home() / ".omnius" / "models" / "asr" / "voxtral-transformers"),
        )
    )
    return root / model_id


def select_device(torch: Any, minimum_vram_gb: float) -> tuple[Any, Any, dict[str, Any]]:
    allow_cpu = os.environ.get("OMNIUS_ASR_ALLOW_CPU", "").lower() in (
        "1",
        "true",
        "yes",
        "on",
    )
    if not torch.cuda.is_available():
        if allow_cpu:
            return torch.device("cpu"), torch.float32, {"device": "cpu"}
        raise RuntimeError(
            "CUDA is required for Voxtral ASR; OMNIUS_ASR_ALLOW_CPU=1 is diagnostic only"
        )
    index = int(os.environ.get("OMNIUS_ASR_CUDA_DEVICE", "0") or "0")
    if index < 0 or index >= torch.cuda.device_count():
        raise RuntimeError(
            f"CUDA device {index} is outside the visible range (count={torch.cuda.device_count()})"
        )
    torch.cuda.set_device(index)
    props = torch.cuda.get_device_properties(index)
    total_gb = props.total_memory / 1024**3
    free_bytes, _ = torch.cuda.mem_get_info(index)
    if total_gb < minimum_vram_gb:
        raise RuntimeError(
            f"{props.name} has {total_gb:.1f} GB VRAM; this Voxtral tier requires approximately {minimum_vram_gb:.0f} GB"
        )
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    device = torch.device(f"cuda:{index}")
    return device, dtype, {
        "device": str(device),
        "name": props.name,
        "freeVramGb": round(free_bytes / 1024**3, 2),
        "totalVramGb": round(total_gb, 2),
        "dtype": str(dtype).replace("torch.", ""),
        "architecture": platform.machine(),
    }


def pull(model_id: str) -> Path:
    from huggingface_hub import snapshot_download

    spec = MODELS[model_id]
    target = model_path(model_id)
    target.mkdir(parents=True, exist_ok=True)
    snapshot_download(repo_id=spec["upstream"], local_dir=str(target))
    (target / ".omnius-model.json").write_text(
        json.dumps(
            {
                "engineId": "voxtral-transformers",
                "modelId": model_id,
                "upstreamModelId": spec["upstream"],
                "pulledAt": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    return target


def load(args: argparse.Namespace) -> tuple[Any, Any, Any, dict[str, Any]]:
    import torch
    from transformers import AutoProcessor

    spec = MODELS[args.model]
    device, dtype, hardware = select_device(torch, spec["minimum_vram_gb"])
    target = model_path(args.model)
    if not (target / "config.json").exists():
        raise RuntimeError(
            f"weights are absent at {target}; call POST /v1/asr/engines/voxtral-transformers/models/{args.model}/pull"
        )
    processor = AutoProcessor.from_pretrained(str(target))
    if spec["realtime"]:
        from transformers import VoxtralRealtimeForConditionalGeneration

        model_class = VoxtralRealtimeForConditionalGeneration
    else:
        from transformers import VoxtralForConditionalGeneration

        model_class = VoxtralForConditionalGeneration
    model = model_class.from_pretrained(
        str(target), torch_dtype=dtype, low_cpu_mem_usage=True
    )
    model.to(device)
    model.eval()
    return processor, model, device, hardware


def transcribe(
    args: argparse.Namespace,
    processor: Any,
    model: Any,
    device: Any,
    audio: Any,
    sample_rate: int,
) -> str:
    import torch

    spec = MODELS[args.model]
    if spec["realtime"]:
        inputs = processor(audio, sampling_rate=sample_rate, return_tensors="pt")
    else:
        request: dict[str, Any] = {
            "audio": audio,
            "model_id": spec["upstream"],
        }
        if args.language:
            request["language"] = args.language
        inputs = processor.apply_transcription_request(**request)
    inputs = {name: value.to(device) for name, value in inputs.items()}
    input_length = inputs["input_ids"].shape[-1] if "input_ids" in inputs else 0
    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=1024)
    if input_length:
        output = output[:, input_length:]
    return processor.batch_decode(output, skip_special_tokens=True)[0].strip()


def transcribe_file(
    args: argparse.Namespace, processor: Any, model: Any, device: Any
) -> int:
    import librosa

    started = time.monotonic()
    sample_rate = int(
        getattr(getattr(processor, "feature_extractor", None), "sampling_rate", 16000)
    )
    audio, _ = librosa.load(args.file, sr=sample_rate, mono=True)
    text = transcribe(args, processor, model, device, audio, sample_rate)
    emit(
        {
            "type": "transcript",
            "text": text,
            "rawText": text,
            "isFinal": True,
            "duration": round(time.monotonic() - started, 3),
            "language": args.language,
            "segments": [],
            "engineId": "voxtral-transformers",
            "modelId": args.model,
        }
    )
    return 0


def stream(
    args: argparse.Namespace, processor: Any, model: Any, device: Any
) -> int:
    import numpy as np

    chunk_bytes = max(2, int(args.chunk_seconds * args.sample_rate * 2))
    window_samples = max(1, int(args.window_seconds * args.sample_rate))
    audio_buffer = np.zeros(0, dtype=np.float32)
    last_text = ""
    emit({"type": "ready"})
    while True:
        data = sys.stdin.buffer.read(chunk_bytes)
        if not data:
            break
        samples = np.frombuffer(data, dtype="<i2").astype(np.float32) / 32768.0
        audio_buffer = np.concatenate((audio_buffer, samples))[-window_samples:]
        if len(audio_buffer) < args.sample_rate:
            continue
        text = transcribe(
            args, processor, model, device, audio_buffer, args.sample_rate
        )
        if text and text != last_text:
            last_text = text
            emit({"type": "transcript", "text": text, "isFinal": False})
    if len(audio_buffer) >= args.sample_rate:
        text = transcribe(
            args, processor, model, device, audio_buffer, args.sample_rate
        )
        if text:
            emit({"type": "transcript", "text": text, "isFinal": True})
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Omnius Voxtral ASR worker")
    parser.add_argument("--model", required=True, choices=sorted(MODELS))
    parser.add_argument("--setup", action="store_true")
    parser.add_argument("--check", action="store_true")
    parser.add_argument("--file")
    parser.add_argument("--language")
    parser.add_argument("--sample-rate", type=int, default=16000)
    parser.add_argument("--chunk-seconds", type=float, default=0.48)
    parser.add_argument("--window-seconds", type=float, default=30.0)
    args = parser.parse_args()
    if args.check:
        emit({"type": "check", "ok": True, "script": str(Path(__file__).resolve())})
        return 0
    if args.setup:
        target = pull(args.model)
        processor, model, device, hardware = load(args)
        del processor, model, device
        emit(
            {
                "type": "ready",
                "engineId": "voxtral-transformers",
                "modelId": args.model,
                "modelPath": str(target),
                "device": hardware["device"],
                "cuda": str(hardware["device"]).startswith("cuda"),
                "hardware": hardware,
            }
        )
        return 0
    processor, model, device, hardware = load(args)
    emit({"type": "status", "message": f"Loaded {args.model} on {hardware['device']}"})
    if args.file:
        return transcribe_file(args, processor, model, device)
    return stream(args, processor, model, device)


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(0)
    except Exception as exc:
        error(f"{type(exc).__name__}: {exc}")
        sys.exit(1)