#!/usr/bin/env python3
"""Detect speech in an audio file and write the regions out as WebVTT.

    python example.py speech.wav
    python example.py speech.mp3 --dtype fp16 --device cpu -o regions.vtt

The VTT is cues covering the detected speech, with no text -- useful for
eyeballing the VAD's decisions in any media player, and as the segmentation
step in front of an ASR model.
"""
from __future__ import annotations

import argparse
from pathlib import Path

import torch

from vad_torch import TorchVadRunner, DEFAULT_WEIGHTS, SAMPLE_RATE


def fmt_vtt_time(seconds: float) -> str:
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}"


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("audio", type=Path)
    p.add_argument("-o", "--out", type=Path, default=None,
                   help="output .vtt (default: alongside the input)")
    p.add_argument("--weights", type=Path, default=DEFAULT_WEIGHTS)
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--dtype", default="fp32", choices=["fp32", "fp16"],
                   help="fp32 (default) reproduces the ONNX reference exactly; "
                        "fp16 is ~2x faster and moves a few boundaries")
    p.add_argument("--batch-size", type=int, default=8)
    p.add_argument("--threshold", type=float, default=0.4,
                   help="in LOGIT space, not probability (default 0.4 ~= 0.60)")
    a = p.parse_args()

    if not a.audio.exists():
        raise SystemExit(f"no such file: {a.audio}")

    try:
        import librosa
    except ImportError:
        raise SystemExit("this example needs librosa: pip install librosa")

    print(f"loading {a.audio.name}")
    audio, _ = librosa.load(str(a.audio), sr=SAMPLE_RATE, mono=True)
    duration = len(audio) / SAMPLE_RATE
    print(f"  {duration:.1f}s at {SAMPLE_RATE} Hz mono")

    dt = {"fp32": torch.float32, "fp16": torch.float16}[a.dtype]
    if a.dtype == "fp16" and a.device == "cpu":
        raise SystemExit("fp16 on CPU is not supported; use --dtype fp32")

    runner = TorchVadRunner(weights=a.weights, device=a.device, dtype=dt,
                            batch_size=a.batch_size)
    print(f"  model on {a.device}, {a.dtype}")

    intervals = runner.speech_intervals(audio, threshold=a.threshold)
    speech = sum(e - s for s, e in intervals)
    print(f"\n{len(intervals)} speech regions, {speech:.1f}s of {duration:.1f}s "
          f"({speech / max(duration, 1e-9) * 100:.1f}% active)")

    out = a.out or a.audio.with_suffix(".vtt")
    lines = ["WEBVTT", ""]
    for i, (start, end) in enumerate(intervals, 1):
        lines.append(str(i))
        lines.append(f"{fmt_vtt_time(start)} --> {fmt_vtt_time(end)}")
        lines.append(f"[speech {i}]")
        lines.append("")
    out.write_text("\n".join(lines), encoding="utf-8")
    print(f"-> {out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())