"""Sentence-level audio segmentation using Silero VAD.

Splits a long audio clip into per-sentence WAVs based on natural speech
boundaries. Tuned for sentences (not word-level chunks) by using a longer
silence-duration threshold to merge intra-sentence pauses but split between
sentences.

Usage:
    python scripts/segment_audio.py INPUT_WAV OUTPUT_DIR

Outputs:
    OUTPUT_DIR/seg_001.wav, seg_002.wav, ...
    OUTPUT_DIR/manifest.json    (start, end, duration, filename per segment)
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import numpy as np
import soundfile as sf
import torch

SAMPLE_RATE = 16000

# Tuned for sentence-level splits (not word-level).
VAD_THRESHOLD = 0.5             # speech probability threshold
MIN_SPEECH_MS = 300             # ignore blips shorter than this
MIN_SILENCE_MS = 500            # silence > this = sentence boundary
SPEECH_PAD_MS = 200             # pad each segment to avoid clipping word edges


def segment(input_wav: Path, out_dir: Path) -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)

    # Load audio
    audio, sr = sf.read(str(input_wav), dtype="float32")
    if sr != SAMPLE_RATE:
        raise ValueError(f"Expected {SAMPLE_RATE} Hz, got {sr}. Re-export with ffmpeg -ar 16000.")
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    audio_t = torch.from_numpy(audio)

    # Load Silero VAD (cached from prior Phase 1 setup)
    model, utils = torch.hub.load(
        repo_or_dir="snakers4/silero-vad",
        model="silero_vad",
        force_reload=False,
        trust_repo=True,
    )
    get_speech_timestamps = utils[0]

    timestamps = get_speech_timestamps(
        audio_t, model,
        sampling_rate=SAMPLE_RATE,
        threshold=VAD_THRESHOLD,
        min_speech_duration_ms=MIN_SPEECH_MS,
        min_silence_duration_ms=MIN_SILENCE_MS,
        speech_pad_ms=SPEECH_PAD_MS,
    )

    # Write each segment as its own WAV + build manifest
    manifest = []
    for i, ts in enumerate(timestamps, start=1):
        start_s = ts["start"] / SAMPLE_RATE
        end_s = ts["end"] / SAMPLE_RATE
        clip = audio[ts["start"]:ts["end"]]
        fname = f"seg_{i:03d}.wav"
        out_path = out_dir / fname
        sf.write(str(out_path), clip, SAMPLE_RATE, subtype="PCM_16")
        manifest.append({
            "filename": fname,
            "index": i,
            "start_seconds": round(start_s, 3),
            "end_seconds": round(end_s, 3),
            "duration_seconds": round(end_s - start_s, 3),
        })

    manifest_data = {
        "source": str(input_wav),
        "sample_rate": SAMPLE_RATE,
        "vad_settings": {
            "threshold": VAD_THRESHOLD,
            "min_speech_ms": MIN_SPEECH_MS,
            "min_silence_ms": MIN_SILENCE_MS,
            "speech_pad_ms": SPEECH_PAD_MS,
        },
        "n_segments": len(manifest),
        "segments": manifest,
    }
    manifest_path = out_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest_data, indent=2), encoding="utf-8")

    return manifest_data


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: segment_audio.py INPUT_WAV OUTPUT_DIR")
        sys.exit(1)
    in_wav = Path(sys.argv[1])
    out_dir = Path(sys.argv[2])
    if not in_wav.exists():
        print(f"Input not found: {in_wav}")
        sys.exit(1)

    result = segment(in_wav, out_dir)
    print(f"Wrote {result['n_segments']} segments to {out_dir}")
    for s in result["segments"]:
        print(f"  {s['filename']}  {s['start_seconds']:6.2f}s -> {s['end_seconds']:6.2f}s  ({s['duration_seconds']:.2f}s)")