"""
SmartFile Speaker Diarization — Hugging Face Space
==================================================
Runs pyannote/speaker-diarization-3.1 to detect WHO SPEAKS WHEN in an audio
clip. SmartFile calls this BEFORE transcription so each speaker turn can be
transcribed and labeled (Locuteur 1, Locuteur 2, ...).

Contract (same simple pattern as the STT Space):
  INPUT  : base64-encoded audio (WAV/etc) in a Textbox.
  OUTPUT : a JSON string -> {"segments": [{"speaker": "...", "start": s, "end": s}, ...],
                             "num_speakers": N}
           or {"error": "..."} on failure.

The pipeline is GATED on Hugging Face and requires accepting terms for BOTH:
  - pyannote/speaker-diarization-3.1
  - pyannote/segmentation-3.0   (dependency)
and an HF_TOKEN secret on this Space with access to them.

Tuning: this client's audio is mostly 2-person interviews/calls, so we hint the
pipeline toward 2 speakers (min 1, max 4) — detected, not hard-forced, so a
stray third voice won't break it. Tiny adjacent same-speaker turns are merged.

NOTE: pyannote loads several models (segmentation + embedding); it is HEAVIER
than the STT Space. Use a paid CPU (or GPU) tier for usable speed.
"""

import base64
import io
import json
import os
import gradio as gr
import torch
from pyannote.audio import Pipeline

HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
MODEL_ID = "pyannote/speaker-diarization-3.1"

# Hint for the typical case (2-person interview/call). Detected, not hard-forced.
DEFAULT_MIN_SPEAKERS = 1
DEFAULT_MAX_SPEAKERS = 4
MERGE_GAP_SEC = 0.5   # merge same-speaker segments separated by < this gap
MIN_SEG_SEC   = 0.3   # drop micro-segments shorter than this (noise/clicks)

print(f"[diar] loading {MODEL_ID} ...")
pipeline = Pipeline.from_pretrained(MODEL_ID, token=HF_TOKEN)
if torch.cuda.is_available():
    pipeline.to(torch.device("cuda"))
    print("[diar] using GPU")
print("[diar] pipeline ready")


def _merge_segments(segs):
    """Merge adjacent same-speaker segments with tiny gaps; drop micro-segments."""
    if not segs:
        return segs
    segs = sorted(segs, key=lambda s: s["start"])
    merged = [segs[0]]
    for s in segs[1:]:
        last = merged[-1]
        if s["speaker"] == last["speaker"] and s["start"] - last["end"] <= MERGE_GAP_SEC:
            last["end"] = max(last["end"], s["end"])
        else:
            merged.append(s)
    return [s for s in merged if (s["end"] - s["start"]) >= MIN_SEG_SEC]


def diarize_b64(audio_b64):
    if not audio_b64:
        return json.dumps({"error": "empty input"})
    try:
        if audio_b64.strip().startswith("data:") and "," in audio_b64[:64]:
            audio_b64 = audio_b64.split(",", 1)[1]
        raw = base64.b64decode(audio_b64)
        print(f"[diar] decoded {len(raw)} audio bytes")

        # pyannote reads a file path or a waveform; easiest robust path is to
        # write the bytes to a temp file (it handles WAV/MP3/etc + resampling).
        import tempfile
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            tmp.write(raw)
            tmp_path = tmp.name

        diarization = pipeline(
            tmp_path,
            min_speakers=DEFAULT_MIN_SPEAKERS,
            max_speakers=DEFAULT_MAX_SPEAKERS,
        )
        os.unlink(tmp_path)

        segs = []
        # pyannote 4.x returns an object whose `.speaker_diarization` yields
        # (turn, speaker) pairs. pyannote 3.x returns an Annotation with
        # `.itertracks(yield_label=True)` yielding (turn, _, speaker). Support
        # both so the Space isn't tied to one pyannote version.
        if hasattr(diarization, "speaker_diarization"):
            for turn, speaker in diarization.speaker_diarization:
                segs.append({
                    "speaker": str(speaker),
                    "start": round(float(turn.start), 2),
                    "end": round(float(turn.end), 2),
                })
        else:
            for turn, _, speaker in diarization.itertracks(yield_label=True):
                segs.append({
                    "speaker": str(speaker),       # e.g. "SPEAKER_00"
                    "start": round(float(turn.start), 2),
                    "end": round(float(turn.end), 2),
                })
        segs = _merge_segments(segs)

        speakers = sorted({s["speaker"] for s in segs})
        print(f"[diar] {len(segs)} segments, {len(speakers)} speakers")
        return json.dumps({"segments": segs, "num_speakers": len(speakers)})
    except Exception as e:
        print(f"[diar] error: {e}")
        return json.dumps({"error": str(e)})


demo = gr.Interface(
    fn=diarize_b64,
    inputs=gr.Textbox(label="Base64 audio"),
    outputs=gr.Textbox(label="Diarization JSON"),
    title="SmartFile Speaker Diarization",
    description="Detect who speaks when (pyannote 3.1). Send base64 audio, get speaker segments as JSON.",
)

if __name__ == "__main__":
    demo.launch()