"""pyannote 3.1 speaker diarization pipeline."""

import os
import warnings
from functools import lru_cache

import huggingface_hub as _hfh
from huggingface_hub.utils import disable_progress_bars
import torch

# pyannote.audio 3.3.x calls hf_hub_download(use_auth_token=...), but
# huggingface_hub >= 0.26 dropped that kwarg in favor of `token`. Patch before
# importing pyannote so its bound reference picks up the shim.
_orig_hf_download = _hfh.hf_hub_download


def _compat_hf_download(*args, **kwargs):
    if "use_auth_token" in kwargs:
        kwargs["token"] = kwargs.pop("use_auth_token")
    return _orig_hf_download(*args, **kwargs)


_hfh.hf_hub_download = _compat_hf_download

# Silence benign pyannote / HF internals that clutter the console:
# - tqdm progress bar from hf_hub_download when pyannote fetches config.yaml
# - ReproducibilityWarning about TF32 being disabled (pyannote does this on purpose)
# - torch UserWarning from stats pooling when a segment has only 1 frame
#   (correction=1 in std() produces dof<=0; output is still well-defined)
disable_progress_bars()
warnings.filterwarnings("ignore", message=r"TensorFloat-32 \(TF32\) has been disabled.*")
warnings.filterwarnings(
    "ignore",
    message=r"std\(\): degrees of freedom is <= 0.*",
    category=UserWarning,
)

from pyannote.audio import Pipeline  # noqa: E402

PYANNOTE_MODEL = "pyannote/speaker-diarization-3.1"


@lru_cache(maxsize=1)
def get_diarization_pipeline():
    token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
    if not token:
        raise RuntimeError(
            "HF_TOKEN env var required. Accept the agreement at "
            "https://huggingface.co/pyannote/speaker-diarization-3.1 and set HF_TOKEN."
        )
    pipe = Pipeline.from_pretrained(PYANNOTE_MODEL, use_auth_token=token)
    if torch.cuda.is_available():
        pipe.to(torch.device("cuda"))
    return pipe


def diarize(audio_path, num_speakers=None):
    """Run diarization. Returns list of {start, end, speaker}."""
    pipe = get_diarization_pipeline()
    kwargs = {}
    if num_speakers and int(num_speakers) > 0:
        kwargs["num_speakers"] = int(num_speakers)
    annotation = pipe(audio_path, **kwargs)
    return [
        {"start": turn.start, "end": turn.end, "speaker": speaker}
        for turn, _, speaker in annotation.itertracks(yield_label=True)
    ]