"""
Local speaker diarization ("who spoke when").

Three steps:
1. sherpa-onnx (pyannote segmentation-3.0, ONNX) cuts the recording into
   single-speaker turns.
2. A neural speaker-embedding model turns each turn into a voice fingerprint.
3. The fingerprints are clustered into speakers. Only turns long enough to
   give a trustworthy fingerprint decide how many speakers there are; short
   interjections ("okay", "so") are then handed to the closest voice.

Everything runs offline on CPU through ONNX Runtime; no Hugging Face token
and no PyTorch.
"""
from __future__ import annotations

import logging
from typing import Callable, List, Optional, Sequence, Tuple

import numpy as np

from config.settings import CPU_THREADS, MAX_AUTO_SPEAKERS, SAME_VOICE_SIMILARITY
from modules.audio_processor import SAMPLE_RATE
from modules.model_downloader import (
    EMBEDDING_MODEL,
    SEGMENTATION_MODEL,
    ensure_diarization_models,
)
from modules.models import DiarizationSegment

logger = logging.getLogger(__name__)

# A turn shorter than this gives an unreliable fingerprint: it may not vote on
# the number of speakers, it is only assigned to one afterwards.
MIN_RELIABLE_SECONDS = 1.5
# Below this the embedding model has too little audio to produce anything.
MIN_EMBED_SECONDS = 0.4
# With auto-detection, a cluster needs this much reliable speech to count as a
# person: 2% of the recording's reliable speech, but never under 4 s or over 20 s.
# A real participant who says less than that is merged into the closest voice;
# choosing the speaker count at upload bypasses this.
MIN_SPEAKER_SECONDS = 4.0
MIN_SPEAKER_SHARE = 0.02
MAX_FLOOR_SECONDS = 20.0


class DiarizationService:
    """Speaker diarization on 16 kHz mono float32 samples."""

    def diarize(
        self,
        audio: np.ndarray,
        num_speakers: Optional[int] = None,
        on_progress: Optional[Callable[[float], None]] = None,
    ) -> List[DiarizationSegment]:
        """
        Return speaker turns ordered by start time.

        *num_speakers* fixes the speaker count when known (most accurate).
        When omitted the count is estimated from how well the voices separate.
        """
        import sherpa_onnx

        ensure_diarization_models()
        embedding_config = sherpa_onnx.SpeakerEmbeddingExtractorConfig(
            model=str(EMBEDDING_MODEL),
            num_threads=CPU_THREADS,
        )

        # -- 1. turns. sherpa-onnx also clusters, with a fixed distance
        # threshold that over-splits badly on real recordings (16 "speakers"
        # for a two-person interview); only its turn boundaries are kept.
        config = sherpa_onnx.OfflineSpeakerDiarizationConfig(
            segmentation=sherpa_onnx.OfflineSpeakerSegmentationModelConfig(
                pyannote=sherpa_onnx.OfflineSpeakerSegmentationPyannoteModelConfig(
                    model=str(SEGMENTATION_MODEL)
                ),
                num_threads=CPU_THREADS,
            ),
            embedding=embedding_config,
            clustering=sherpa_onnx.FastClusteringConfig(num_clusters=-1, threshold=0.5),
            # Ignore blips shorter than 0.3 s; bridge pauses shorter than 0.5 s.
            min_duration_on=0.3,
            min_duration_off=0.5,
        )
        if not config.validate():
            raise RuntimeError("Diarization models are missing or invalid.")

        def _callback(done: int, total: int) -> int:
            if on_progress and total:
                on_progress(0.7 * done / total)
            return 0

        result = sherpa_onnx.OfflineSpeakerDiarization(config).process(
            audio, callback=_callback
        ).sort_by_start_time()
        spans = [(r.start, r.end) for r in result]
        if not spans:
            return []

        # -- 2. one fingerprint per turn
        extractor = sherpa_onnx.SpeakerEmbeddingExtractor(embedding_config)
        embeddings: list[Optional[np.ndarray]] = []
        for i, (start, end) in enumerate(spans):
            embeddings.append(
                _embed(extractor, audio[int(start * SAMPLE_RATE): int(end * SAMPLE_RATE)])
            )
            if on_progress:
                on_progress(0.7 + 0.3 * (i + 1) / len(spans))

        # -- 3. speakers
        labels = cluster_speakers(
            [end - start for start, end in spans], embeddings, num_speakers
        )
        turns = label_turns([(s, e, label) for (s, e), label in zip(spans, labels)])

        talk_time: dict[str, float] = {}
        for t in turns:
            talk_time[t.speaker] = talk_time.get(t.speaker, 0.0) + t.end - t.start
        logger.info(
            "Diarization complete: %d turns, %d speakers (%s)",
            len(turns), len(talk_time),
            ", ".join(f"{name} {seconds:.0f}s" for name, seconds in talk_time.items()),
        )
        return turns


def _embed(extractor, samples: np.ndarray) -> Optional[np.ndarray]:
    """Unit-length voice fingerprint for *samples*, or None if there is too little audio."""
    if len(samples) < MIN_EMBED_SECONDS * SAMPLE_RATE:
        return None
    stream = extractor.create_stream()
    stream.accept_waveform(sample_rate=SAMPLE_RATE, waveform=samples)
    stream.input_finished()
    if not extractor.is_ready(stream):
        return None
    embedding = np.asarray(extractor.compute(stream), dtype=np.float32)
    norm = np.linalg.norm(embedding)
    return embedding / norm if norm > 0 else None


def cluster_speakers(
    durations: Sequence[float],
    embeddings: Sequence[Optional[np.ndarray]],
    num_speakers: Optional[int] = None,
) -> List[int]:
    """
    Return a speaker index for every turn.

    Reliable turns (long enough for a trustworthy fingerprint) are clustered;
    every other turn joins the closest cluster, or, with no fingerprint at
    all, the speaker of the turn before it.
    """
    from sklearn.cluster import AgglomerativeClustering

    n = len(durations)
    reliable = [
        i for i in range(n)
        if embeddings[i] is not None and durations[i] >= MIN_RELIABLE_SECONDS
    ]
    if len(reliable) < 2:  # very short recording: use whatever fingerprints exist
        reliable = [i for i in range(n) if embeddings[i] is not None]
    if len(reliable) < 2:
        return [0] * n

    fixed = bool(num_speakers and num_speakers > 0)
    while True:
        X = np.stack([embeddings[i] for i in reliable])
        k = min(num_speakers if fixed else estimate_speaker_count(X), len(reliable))
        if k <= 1:
            return [0] * n
        reliable_labels = AgglomerativeClustering(
            n_clusters=k, metric="cosine", linkage="average"
        ).fit_predict(X)
        if fixed:
            break
        # A cluster holding almost no speech is an odd-sounding turn (laughter,
        # applause, a change of microphone), not a person. Clustering isolates
        # such turns eagerly, so set them aside and count again without them.
        seconds = np.array([durations[i] for i in reliable])
        floor = min(MAX_FLOOR_SECONDS, max(MIN_SPEAKER_SECONDS, MIN_SPEAKER_SHARE * seconds.sum()))
        minor = [c for c in range(k) if seconds[reliable_labels == c].sum() < floor]
        if not minor:
            break
        reliable = [i for i, label in zip(reliable, reliable_labels) if label not in minor]
        if len(reliable) < 2:
            return [0] * n

    if not fixed:
        reliable_labels = merge_same_voices(X, reliable_labels)
    k = int(reliable_labels.max()) + 1
    if k <= 1:
        return [0] * n
    centroids = _centroids(X, reliable_labels)

    labels: list[int] = []
    for i in range(n):
        if embeddings[i] is None:
            labels.append(labels[-1] if labels else -1)
        else:
            labels.append(int(np.argmax(centroids @ embeddings[i])))
    for i, label in zip(reliable, reliable_labels):
        labels[i] = int(label)
    # Leading turns without a fingerprint take the first known speaker.
    first_known = next(label for label in labels if label != -1)
    return [first_known if label == -1 else label for label in labels]


def _centroids(X: np.ndarray, labels: np.ndarray) -> np.ndarray:
    """Average voice of each cluster, unit length, indexed by label."""
    centroids = np.stack([X[labels == c].mean(axis=0) for c in range(int(labels.max()) + 1)])
    return centroids / np.linalg.norm(centroids, axis=1, keepdims=True)


def merge_same_voices(X: np.ndarray, labels: np.ndarray) -> np.ndarray:
    """
    Merge clusters whose average voices are the same person, and renumber.

    The silhouette score always finds some way to split a recording, even one
    with a single speaker. What tells the cases apart is how alike the
    resulting voices are: on test recordings, different people scored a
    similarity of 0.24 to 0.39 and one person split in two scored 0.92.
    """
    labels = np.unique(labels, return_inverse=True)[1]
    while labels.max() > 0:
        similarity = _centroids(X, labels) @ _centroids(X, labels).T
        np.fill_diagonal(similarity, -1.0)
        a, b = np.unravel_index(np.argmax(similarity), similarity.shape)
        logger.info("Most alike pair of voices has similarity %.2f", similarity[a, b])
        if similarity[a, b] < SAME_VOICE_SIMILARITY:
            break
        labels[labels == b] = a
        labels = np.unique(labels, return_inverse=True)[1]
    return labels


def estimate_speaker_count(X: np.ndarray) -> int:
    """
    Pick the number of clusters whose members sit closest to their own cluster
    and furthest from the others (best silhouette score). This never answers
    "one"; merge_same_voices() makes that call afterwards.
    """
    from sklearn.cluster import AgglomerativeClustering
    from sklearn.metrics import silhouette_score

    scores: dict[int, float] = {}
    for k in range(2, min(MAX_AUTO_SPEAKERS, len(X) - 1) + 1):
        labels = AgglomerativeClustering(
            n_clusters=k, metric="cosine", linkage="average"
        ).fit_predict(X)
        scores[k] = float(silhouette_score(X, labels, metric="cosine"))
    logger.info(
        "Speaker count scores from %d reliable turns: %s",
        len(X), ", ".join(f"{k}={s:.2f}" for k, s in scores.items()),
    )
    return max(scores, key=scores.get) if scores else 1


def label_turns(raw_turns: List[Tuple[float, float, int]]) -> List[DiarizationSegment]:
    """Name clusters 'Speaker 1', 'Speaker 2', ... in order of first appearance."""
    numbers: dict[int, int] = {}
    turns: list[DiarizationSegment] = []
    for start, end, cluster in raw_turns:
        number = numbers.setdefault(cluster, len(numbers) + 1)
        turns.append(DiarizationSegment(start=start, end=end, speaker=f"Speaker {number}"))
    return turns