#!/usr/bin/env python3
"""
Transcribe an audio file using openai-whisper (CUDA-only by default).

stdin:  JSON {"path": "...", "model": "tiny|base|small|medium|large-v3"}
stdout: JSON {"text": "...", "duration": null, "segments": []}
        or   {"error": "<kind>", "message": "..."}

ASR is GPU-first: CPU fallback is refused unless OMNIUS_ASR_ALLOW_CPU=1.
"""
import sys, os, json, math, struct, warnings, wave

warnings.filterwarnings("ignore")
os.environ.setdefault("PYTHONWARNINGS", "ignore")

PCM16_SILENCE_MAX_PEAK = 4
PCM16_SILENCE_MAX_RMS = 1.0
PCM16_ACTIVE_SAMPLE_THRESHOLD = 8


def exact_pcm16_silence(path):
    """Return typed no-speech evidence for exact PCM16/16k/mono WAV only.

    This is intentionally not VAD. It reads user bytes but never rewrites,
    normalizes, decodes another media format, or instantiates Whisper.
    """
    try:
        with wave.open(path, "rb") as wav:
            if (
                wav.getcomptype() != "NONE"
                or wav.getnchannels() != 1
                or wav.getsampwidth() != 2
                or wav.getframerate() != 16000
            ):
                return None
            frames = wav.readframes(wav.getnframes())
    except (wave.Error, OSError, EOFError):
        return None
    sample_count = len(frames) // 2
    if sample_count == 0:
        metrics = {
            "sampleRateHz": 16000, "channels": 1, "bitsPerSample": 16,
            "sampleCount": 0, "durationMs": 0.0, "peakPcm16": 0,
            "rmsPcm16": 0.0, "activeSampleCount": 0, "activeSampleRatio": 0.0,
        }
    else:
        samples = struct.iter_unpack("<h", frames[: sample_count * 2])
        total_squares = 0
        peak = 0
        active = 0
        for (sample,) in samples:
            magnitude = abs(sample)
            total_squares += sample * sample
            peak = max(peak, magnitude)
            if magnitude >= PCM16_ACTIVE_SAMPLE_THRESHOLD:
                active += 1
        metrics = {
            "sampleRateHz": 16000, "channels": 1, "bitsPerSample": 16,
            "sampleCount": sample_count, "durationMs": sample_count / 16.0,
            "peakPcm16": peak, "rmsPcm16": math.sqrt(total_squares / sample_count),
            "activeSampleCount": active, "activeSampleRatio": active / sample_count,
        }
    if (
        metrics["peakPcm16"] <= PCM16_SILENCE_MAX_PEAK
        and metrics["rmsPcm16"] <= PCM16_SILENCE_MAX_RMS
        and metrics["activeSampleCount"] == 0
    ):
        return {"kind": "no_speech", "reason": "digital_silence", "text": "", "signal": metrics}
    return None


def main() -> None:
    raw = sys.stdin.read()
    try:
        data = json.loads(raw) if raw.strip() else {}
    except Exception as e:
        print(json.dumps({"error": "bad_input", "message": str(e)}))
        return

    path = data.get("path")
    if not path or not os.path.exists(path):
        print(json.dumps({"error": "file_not_found", "path": path or ""}))
        return

    no_speech = exact_pcm16_silence(path)
    if no_speech is not None:
        print(json.dumps({**no_speech, "duration": no_speech["signal"]["durationMs"] / 1000.0, "segments": []}))
        return

    model_name = data.get("model") or "tiny"
    try:
        import whisper
        device = os.environ.get("OMNIUS_ASR_DEVICE", "cuda").strip().lower() or "cuda"
        allow_cpu = os.environ.get("OMNIUS_ASR_ALLOW_CPU", "").strip().lower() in ("1", "true", "yes", "on")
        if device == "auto":
            device = "cuda"
        if device.startswith("cuda"):
            import torch
            if not torch.cuda.is_available() or torch.cuda.device_count() <= 0:
                raise RuntimeError(
                    f"CUDA-only ASR requested but torch cannot use CUDA "
                    f"(torch.version.cuda={getattr(torch.version, 'cuda', None)}, device_count={torch.cuda.device_count()}). "
                    "Install a CUDA-enabled PyTorch build for this device, or set OMNIUS_ASR_ALLOW_CPU=1 only for emergency fallback."
                )
        elif device == "cpu" and not allow_cpu:
            raise RuntimeError("CPU ASR refused. Set OMNIUS_ASR_ALLOW_CPU=1 only for explicit emergency fallback.")
        model = whisper.load_model(model_name, device=device)
        result = model.transcribe(path, fp16=device.startswith("cuda"), condition_on_previous_text=False)
    except Exception as e:
        print(json.dumps({"error": "whisper_failed", "message": str(e)}))
        return

    text = (result.get("text") or "").strip()
    raw_segments = result.get("segments") or []
    segments = [
        {
            "start": float(s.get("start", 0.0)),
            "end": float(s.get("end", 0.0)),
            "text": (s.get("text") or "").strip(),
        }
        for s in raw_segments
    ]
    duration = float(segments[-1]["end"]) if segments else None
    print(json.dumps({"text": text, "duration": duration, "segments": segments}))


if __name__ == "__main__":
    main()