#!/usr/bin/env python3
"""
whisper-transcribe.py  —  faster-whisper + resemblyzer speaker diarization
Usage:  python whisper-transcribe.py <audio_file> [model_size]
Output: "SPEAKER_X: text" lines on stdout (falls back to plain text if diarization fails)
"""
import sys
import os
import warnings
warnings.filterwarnings("ignore")
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"

SAMPLE_RATE = 16000


def load_audio(path):
    """Load any audio file as 16 kHz mono float32 numpy array."""
    import numpy as np

    ext = os.path.splitext(path)[1].lower()

    # WAV: read directly — no extra dependencies
    if ext == ".wav":
        import wave
        with wave.open(path, "rb") as wf:
            ch  = wf.getnchannels()
            sw  = wf.getsampwidth()
            sr  = wf.getframerate()
            raw = wf.readframes(wf.getnframes())

        if sw == 2:
            audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
        elif sw == 4:
            audio = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / 2147483648.0
        else:
            raise ValueError(f"Unsupported sample width {sw}")

        if ch > 1:                           # stereo -> mono
            audio = audio.reshape(-1, ch).mean(axis=1)

        if sr != SAMPLE_RATE:
            from scipy.signal import resample_poly
            from math import gcd
            g = gcd(SAMPLE_RATE, sr)
            audio = resample_poly(audio, SAMPLE_RATE // g, sr // g)

        return audio.astype(np.float32)

    # Everything else: try soundfile first, then librosa
    try:
        import soundfile as sf
        audio, sr = sf.read(path, dtype="float32", always_2d=False)
        if audio.ndim > 1:
            audio = audio.mean(axis=1)
    except Exception:
        import librosa
        audio, sr = librosa.load(path, sr=None, mono=True, dtype=np.float32)

    if sr != SAMPLE_RATE:
        from scipy.signal import resample_poly
        from math import gcd
        g = gcd(SAMPLE_RATE, sr)
        audio = resample_poly(audio, SAMPLE_RATE // g, sr // g)

    return audio.astype(np.float32)


def embed_segments(audio, segments):
    """Return (embeddings, valid_mask) using resemblyzer VoiceEncoder."""
    import numpy as np
    from resemblyzer import VoiceEncoder

    encoder = VoiceEncoder()
    embeddings  = []
    valid_mask  = []
    min_samples = int(0.4 * SAMPLE_RATE)   # skip segments shorter than 400 ms

    for seg in segments:
        start = int(seg.start * SAMPLE_RATE)
        end   = int(seg.end   * SAMPLE_RATE)
        chunk = audio[start:end]
        if len(chunk) < min_samples:
            embeddings.append(None)
            valid_mask.append(False)
            continue
        try:
            emb = encoder.embed_utterance(chunk)
            embeddings.append(emb)
            valid_mask.append(True)
        except Exception:
            embeddings.append(None)
            valid_mask.append(False)

    return embeddings, valid_mask


def cluster_speakers(embeddings, valid_mask):
    """
    Agglomerative clustering on cosine-normalised embeddings.
    Returns a list of speaker labels (0-based ints) for every segment.
    Segments that failed embedding inherit their nearest neighbour's label.
    """
    import numpy as np
    from sklearn.preprocessing import normalize
    from sklearn.cluster import AgglomerativeClustering

    valid_embs = [e for e, v in zip(embeddings, valid_mask) if v]
    valid_idxs = [i for i, v in enumerate(valid_mask)       if v]

    if not valid_embs:
        return [0] * len(embeddings)

    if len(valid_embs) == 1:
        return [0] * len(embeddings)

    mat = normalize(np.array(valid_embs))
    n   = len(valid_embs)

    # Pick best k (2-5 speakers) via silhouette score
    best_k, best_score = 1, -1.0
    for k in range(2, min(6, n + 1)):
        try:
            from sklearn.metrics import silhouette_score
            cl    = AgglomerativeClustering(n_clusters=k, metric="cosine", linkage="average")
            lbls  = cl.fit_predict(mat)
            score = silhouette_score(mat, lbls, metric="cosine")
            if score > best_score:
                best_score, best_k = score, k
        except Exception:
            pass

    clustering  = AgglomerativeClustering(n_clusters=best_k, metric="cosine", linkage="average")
    cluster_ids = clustering.fit_predict(mat)

    # Map back to all segments
    full_labels = [None] * len(embeddings)
    for pos, idx in enumerate(valid_idxs):
        full_labels[idx] = int(cluster_ids[pos])

    # Fill gaps (short/silent segments) from nearest neighbour
    for i in range(len(full_labels)):
        if full_labels[i] is not None:
            continue
        for delta in range(1, len(full_labels)):
            if i - delta >= 0 and full_labels[i - delta] is not None:
                full_labels[i] = full_labels[i - delta]
                break
            if i + delta < len(full_labels) and full_labels[i + delta] is not None:
                full_labels[i] = full_labels[i + delta]
                break
        if full_labels[i] is None:
            full_labels[i] = 0

    return full_labels


def format_output(segments, speaker_labels):
    """
    Collapse consecutive same-speaker segments then print:
        SPEAKER_1: ...text...
    """
    lines  = []
    buf_sp = None
    buf_tx = []

    for seg, sp_id in zip(segments, speaker_labels):
        text  = seg.text.strip()
        if not text:
            continue
        label = f"SPEAKER_{sp_id + 1}"
        if label == buf_sp:
            buf_tx.append(text)
        else:
            if buf_sp is not None:
                lines.append(f"{buf_sp}: {' '.join(buf_tx)}")
            buf_sp = label
            buf_tx = [text]

    if buf_sp is not None:
        lines.append(f"{buf_sp}: {' '.join(buf_tx)}")

    return lines


def main():
    if len(sys.argv) < 2:
        print("Usage: python whisper-transcribe.py <audio_file> [tiny|base|small|medium]",
              file=sys.stderr)
        sys.exit(1)

    audio_file = sys.argv[1]
    model_size = sys.argv[2] if len(sys.argv) > 2 else "small"

    if not os.path.exists(audio_file):
        print(f"File not found: {audio_file}", file=sys.stderr)
        sys.exit(1)

    try:
        from faster_whisper import WhisperModel

        model = WhisperModel(model_size, device="cpu", compute_type="int8")
        raw_segs, _ = model.transcribe(
            audio_file, beam_size=5, language="en",
            vad_filter=True,        # skip silence automatically
        )
        segments = list(raw_segs)   # materialise (generator is single-pass)

        if not segments:
            sys.stdout.flush()
            return

        # --- Speaker diarization ---
        try:
            audio                  = load_audio(audio_file)
            embeddings, valid_mask = embed_segments(audio, segments)
            speaker_labels         = cluster_speakers(embeddings, valid_mask)
            lines                  = format_output(segments, speaker_labels)
        except Exception as e:
            print(f"[diarization unavailable: {e}]", file=sys.stderr)
            lines = [seg.text.strip() for seg in segments if seg.text.strip()]

        for line in lines:
            print(line)
        sys.stdout.flush()

    except Exception as e:
        print(f"Transcription error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()
