//! Speaker diarization composed from existing in-crate components.
//!
//! Pipeline (~the pyannote-community-1 shape, simpler):
//!   1. Silero V5 VAD → per-32ms speech probability → speech intervals
//!   2. WavLM-SV → 512-d x-vector per ~2-second window inside each interval
//!   3. Agglomerative average-linkage clustering on cosine distance →
//!      speaker labels for every window
//!   4. Merge adjacent same-speaker windows back into time segments
//!
//! Pure candle, zero new deps. Uses [`crate::silero_vad::SileroVad`] (Phase 11)
//! and [`crate::wavlm_sv::WavLmSv`] (Phase 5d). Run as a sidecar binary or
//! library call during data prep — diarization is offline work that doesn't
//! belong in the live voice loop.
//!
//! Reference: ~80% of the WhisperX / pyannote pipeline without the ort
//! runtime risk (see `docs/sesame_gap_analysis.md` and the runtime-conflict
//! notes in `personal_voice_training_guide.md`).

use crate::error::{CsmError, Result};
use crate::silero_vad::SileroVad;
use crate::wavlm_sv::WavLmSv;
use candle_core::Device;

/// Sample rate the diarizer expects on its input. Both Silero V5 and
/// WavLM-SV in this crate operate at 16 kHz; the caller is responsible for
/// resampling 24 kHz CSM-side audio down before calling [`Diarizer::diarize`].
pub const SAMPLE_RATE: u32 = 16_000;

/// Silero V5 chunk = 512 samples = 32 ms at 16 kHz. Mirror the constant
/// from `silero_vad.rs` so we don't recompute it.
const VAD_CHUNK_MS: f32 = 32.0;

/// One speaker-attributed time interval.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DiarizedSegment {
    pub start_s: f32,
    pub end_s: f32,
    pub speaker: usize,
}

/// Configuration for the full pipeline.
#[derive(Debug, Clone)]
pub struct DiarizationConfig {
    /// VAD probability above which a 32 ms chunk is considered speech.
    /// Silero V5 is well-calibrated; 0.5 is a safe default.
    pub vad_threshold: f32,
    /// Minimum continuous speech length (seconds) for a region to be kept.
    /// Filters out cough / lip-smack / single-word interjections that
    /// would produce noisy embeddings.
    pub min_speech_s: f32,
    /// Maximum silence (seconds) inside a speech interval before it is split.
    /// Tolerates short pauses inside a single utterance.
    pub max_silence_s: f32,
    /// WavLM-SV window length in seconds. 2 s is the sweet spot —
    /// long enough for a discriminative x-vector, short enough that
    /// time resolution survives.
    pub window_s: f32,
    /// Hop between consecutive WavLM-SV windows. 1 s = 50 % overlap.
    pub hop_s: f32,
    /// Cosine-distance threshold above which two clusters are NOT merged.
    /// Lower = more clusters (over-segmentation); higher = fewer clusters
    /// (collapsed speakers). 0.5 is a literature default for WavLM-SV.
    pub cluster_threshold: f32,
    /// When `Some(k)`, force exactly `k` clusters and ignore the threshold.
    /// `None` = auto-detect via the threshold.
    pub n_speakers: Option<usize>,
    /// Drop output segments shorter than this (after merge). Removes
    /// 1-window blips that survived clustering.
    pub min_segment_s: f32,
}

impl Default for DiarizationConfig {
    fn default() -> Self {
        Self {
            vad_threshold: 0.5,
            min_speech_s: 0.5,
            max_silence_s: 0.3,
            window_s: 2.0,
            hop_s: 1.0,
            cluster_threshold: 0.5,
            n_speakers: None,
            min_segment_s: 0.5,
        }
    }
}

/// Owns the two underlying models. Construct once, reuse across many calls.
pub struct Diarizer {
    pub vad: SileroVad,
    pub sv: WavLmSv,
    device: Device,
}

impl Diarizer {
    pub fn new(vad: SileroVad, sv: WavLmSv, device: Device) -> Self {
        Self { vad, sv, device }
    }

    /// Run the full pipeline. `samples_16k` must be 16 kHz mono f32.
    pub fn diarize(
        &mut self,
        samples_16k: &[f32],
        cfg: &DiarizationConfig,
    ) -> Result<Vec<DiarizedSegment>> {
        // 1. VAD → speech intervals
        let probs = self
            .vad
            .forward_audio(samples_16k, &self.device)
            .map_err(|e| CsmError::Config(format!("diarize: vad: {e}")))?;
        let intervals = vad_intervals(&probs, cfg);
        if intervals.is_empty() {
            return Ok(Vec::new());
        }
        tracing::debug!("diarize: {} speech intervals", intervals.len());

        // 2. WavLM-SV embeddings per window inside each interval. VAD
        //    interval bounds are chunk-quantized and may exceed
        //    `samples_16k.len()`; clamp before slicing.
        let mut windows: Vec<WindowEmb> = Vec::new();
        let win_samples = (cfg.window_s * SAMPLE_RATE as f32) as usize;
        let hop_samples = (cfg.hop_s * SAMPLE_RATE as f32) as usize;
        for &(s_idx, e_idx) in intervals.iter() {
            let s = s_idx.min(samples_16k.len());
            let e = e_idx.min(samples_16k.len());
            if e <= s {
                continue;
            }
            // Interval shorter than one window: take the whole interval as a
            // single sub-window if it's at least half a window long. This is
            // the common case for short utterances.
            if e - s < win_samples {
                if e - s >= win_samples / 2 {
                    let chunk = &samples_16k[s..e];
                    let emb = self.sv.embed_samples(chunk, &self.device)?;
                    let l2 = l2_normalize(&emb);
                    windows.push(WindowEmb {
                        start_s: s as f32 / SAMPLE_RATE as f32,
                        end_s: e as f32 / SAMPLE_RATE as f32,
                        emb: l2,
                    });
                }
                continue;
            }
            let mut t = s;
            while t + win_samples <= e {
                let chunk = &samples_16k[t..t + win_samples];
                let emb = self.sv.embed_samples(chunk, &self.device)?;
                let l2 = l2_normalize(&emb);
                windows.push(WindowEmb {
                    start_s: t as f32 / SAMPLE_RATE as f32,
                    end_s: (t + win_samples) as f32 / SAMPLE_RATE as f32,
                    emb: l2,
                });
                t += hop_samples;
            }
            // Tail: grab one last full-width window flush against `e` to
            // capture the trailing speech if the hop didn't already cover it.
            if t < e && e - t >= win_samples / 2 {
                let start = e.saturating_sub(win_samples);
                let chunk = &samples_16k[start..e];
                let emb = self.sv.embed_samples(chunk, &self.device)?;
                let l2 = l2_normalize(&emb);
                windows.push(WindowEmb {
                    start_s: start as f32 / SAMPLE_RATE as f32,
                    end_s: e as f32 / SAMPLE_RATE as f32,
                    emb: l2,
                });
            }
        }
        if windows.is_empty() {
            return Ok(Vec::new());
        }
        tracing::debug!("diarize: {} embedding windows", windows.len());

        // 3. Agglomerative clustering
        let labels = agglomerative_cluster(
            &windows.iter().map(|w| w.emb.as_slice()).collect::<Vec<_>>(),
            cfg.cluster_threshold,
            cfg.n_speakers,
        );

        // 4. Merge adjacent same-speaker windows
        let mut segments = Vec::new();
        let mut cur_speaker = labels[0];
        let mut cur_start = windows[0].start_s;
        let mut cur_end = windows[0].end_s;
        for (i, w) in windows.iter().enumerate().skip(1) {
            if labels[i] == cur_speaker && w.start_s <= cur_end + 0.01 {
                cur_end = w.end_s.max(cur_end);
            } else {
                if cur_end - cur_start >= cfg.min_segment_s {
                    segments.push(DiarizedSegment {
                        start_s: cur_start,
                        end_s: cur_end,
                        speaker: cur_speaker,
                    });
                }
                cur_speaker = labels[i];
                cur_start = w.start_s;
                cur_end = w.end_s;
            }
        }
        if cur_end - cur_start >= cfg.min_segment_s {
            segments.push(DiarizedSegment {
                start_s: cur_start,
                end_s: cur_end,
                speaker: cur_speaker,
            });
        }
        Ok(segments)
    }
}

struct WindowEmb {
    start_s: f32,
    end_s: f32,
    emb: Vec<f32>,
}

/// Convert per-chunk Silero V5 probabilities into sample-index intervals
/// `(start, end)`. Internally chunk-quantized; the returned bounds are
/// exact-multiple-of-`CHUNK_SAMPLES`.
fn vad_intervals(probs: &[f32], cfg: &DiarizationConfig) -> Vec<(usize, usize)> {
    let chunk_samples = (VAD_CHUNK_MS * SAMPLE_RATE as f32 / 1000.0) as usize;
    let max_silence_chunks = (cfg.max_silence_s * 1000.0 / VAD_CHUNK_MS).round() as usize;
    let min_speech_chunks = (cfg.min_speech_s * 1000.0 / VAD_CHUNK_MS).ceil() as usize;

    // Smooth small gaps: turn isolated low-prob chunks inside speech
    // into speech if surrounding `max_silence_chunks` is mostly speech.
    let mut active: Vec<bool> = probs.iter().map(|p| *p >= cfg.vad_threshold).collect();
    let mut last_speech: Option<usize> = None;
    for i in 0..active.len() {
        if active[i] {
            if let Some(prev) = last_speech
                && i - prev <= max_silence_chunks + 1
            {
                active[(prev + 1)..i].fill(true);
            }
            last_speech = Some(i);
        }
    }

    // Coalesce contiguous true runs into intervals.
    let mut intervals = Vec::new();
    let mut i = 0;
    while i < active.len() {
        if active[i] {
            let start = i;
            while i < active.len() && active[i] {
                i += 1;
            }
            let len = i - start;
            if len >= min_speech_chunks {
                intervals.push((start * chunk_samples, i * chunk_samples));
            }
        } else {
            i += 1;
        }
    }
    intervals
}

fn l2_normalize(v: &[f32]) -> Vec<f32> {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-9);
    v.iter().map(|x| x / norm).collect()
}

/// Cosine *distance* between two L2-normalized vectors: `1 - dot(a, b)`.
/// Range `[0, 2]`; identical = 0; orthogonal = 1; opposite = 2.
fn cosine_distance_normed(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    1.0 - dot
}

/// Agglomerative average-linkage clustering on cosine distance.
///
/// Each input vector is a starting cluster. At each step, find the two
/// clusters with the smallest pairwise average distance and merge them.
/// Stop when:
///   - `n_speakers` is `Some(k)` and we have `k` clusters, OR
///   - the smallest distance exceeds `threshold` (auto-detect mode)
///
/// Returns one label per input vector. Labels are remapped to be a dense
/// `0..n_clusters` range in order of first appearance.
fn agglomerative_cluster(embs: &[&[f32]], threshold: f32, n_speakers: Option<usize>) -> Vec<usize> {
    let n = embs.len();
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![0];
    }

    // Each cluster carries: its membership (Vec of input idx) and the
    // current label. We work with cluster ids 0..n that get unioned.
    let mut clusters: Vec<Vec<usize>> = (0..n).map(|i| vec![i]).collect();

    // Pairwise distance cache, keyed by sorted (smaller_id, larger_id).
    // We rebuild after each merge — n is small (~ 10²) for typical
    // diarization, so O(n²) per step is fine.
    loop {
        if clusters.len() == 1 {
            break;
        }
        if let Some(k) = n_speakers
            && clusters.len() <= k
        {
            break;
        }
        // Find the closest pair using average linkage.
        let mut best: Option<(usize, usize, f32)> = None;
        for i in 0..clusters.len() {
            for j in (i + 1)..clusters.len() {
                let mut sum = 0.0f32;
                let mut count = 0u32;
                for &a in &clusters[i] {
                    for &b in &clusters[j] {
                        sum += cosine_distance_normed(embs[a], embs[b]);
                        count += 1;
                    }
                }
                let d = sum / count as f32;
                if best.map(|(_, _, bd)| d < bd).unwrap_or(true) {
                    best = Some((i, j, d));
                }
            }
        }
        let (i, j, d) = best.unwrap();
        if n_speakers.is_none() && d > threshold {
            break;
        }
        // Merge j into i, remove j (j > i so popping order is safe).
        let merged_j = clusters.remove(j);
        clusters[i].extend(merged_j);
    }

    // Build per-input label mapping. Remap to dense 0..k in order of
    // first appearance in the input.
    let mut raw_label = vec![0usize; n];
    for (cid, members) in clusters.iter().enumerate() {
        for &m in members {
            raw_label[m] = cid;
        }
    }
    let mut remap = std::collections::HashMap::new();
    let mut next_label = 0usize;
    let mut out = Vec::with_capacity(n);
    for &r in &raw_label {
        let lab = *remap.entry(r).or_insert_with(|| {
            let l = next_label;
            next_label += 1;
            l
        });
        out.push(lab);
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cosine_distance_identical_is_zero() {
        let a = l2_normalize(&[1.0, 2.0, 3.0]);
        let d = cosine_distance_normed(&a, &a);
        assert!(d.abs() < 1e-6, "identical vectors → distance ≈ 0, got {d}");
    }

    #[test]
    fn cosine_distance_orthogonal_is_one() {
        let a = l2_normalize(&[1.0, 0.0]);
        let b = l2_normalize(&[0.0, 1.0]);
        let d = cosine_distance_normed(&a, &b);
        assert!((d - 1.0).abs() < 1e-6, "orthogonal → distance ≈ 1, got {d}");
    }

    #[test]
    fn agglomerative_two_clusters_merges_within_threshold() {
        // Two well-separated clusters of 3 vectors each.
        let cluster_a = [
            l2_normalize(&[1.0, 0.0, 0.0]),
            l2_normalize(&[0.95, 0.1, 0.0]),
            l2_normalize(&[0.9, 0.0, 0.1]),
        ];
        let cluster_b = [
            l2_normalize(&[0.0, 1.0, 0.0]),
            l2_normalize(&[0.0, 0.95, 0.1]),
            l2_normalize(&[0.05, 0.9, 0.0]),
        ];
        let mut owned = Vec::new();
        owned.extend(cluster_a.iter().cloned());
        owned.extend(cluster_b.iter().cloned());
        let refs: Vec<&[f32]> = owned.iter().map(|v| v.as_slice()).collect();
        let labels = agglomerative_cluster(&refs, 0.5, None);
        // First three should share a label; last three should share the other.
        assert_eq!(labels[0], labels[1]);
        assert_eq!(labels[0], labels[2]);
        assert_eq!(labels[3], labels[4]);
        assert_eq!(labels[3], labels[5]);
        assert_ne!(labels[0], labels[3]);
    }

    #[test]
    fn agglomerative_n_speakers_forces_k_clusters() {
        // 4 well-separated points in different directions; threshold-mode
        // would give 4 clusters, but n_speakers=2 should collapse them
        // to exactly 2.
        let pts = [
            l2_normalize(&[1.0, 0.0, 0.0, 0.0]),
            l2_normalize(&[0.0, 1.0, 0.0, 0.0]),
            l2_normalize(&[0.0, 0.0, 1.0, 0.0]),
            l2_normalize(&[0.0, 0.0, 0.0, 1.0]),
        ];
        let refs: Vec<&[f32]> = pts.iter().map(|v| v.as_slice()).collect();
        let labels = agglomerative_cluster(&refs, 0.1, Some(2));
        let unique: std::collections::HashSet<_> = labels.iter().collect();
        assert_eq!(unique.len(), 2, "n_speakers=2 should yield 2 clusters");
    }

    #[test]
    fn vad_intervals_drops_too_short_speech() {
        // 32 ms chunks, threshold 0.5. Pattern: 0.1 0.9 0.1 0.1 0.9 0.9 0.9 0.9 0.1
        // = one isolated speech chunk (skipped, below min_speech) then 4
        // contiguous (kept). With min_speech_s=0.1 (≈3 chunks at 32ms).
        let probs = [0.1, 0.9, 0.1, 0.1, 0.9, 0.9, 0.9, 0.9, 0.1];
        let cfg = DiarizationConfig {
            vad_threshold: 0.5,
            min_speech_s: 0.1,
            max_silence_s: 0.0,
            ..DiarizationConfig::default()
        };
        let intervals = vad_intervals(&probs, &cfg);
        assert_eq!(intervals.len(), 1, "should drop the isolated speech chunk");
        // 4 chunks × 512 samples = 2048 samples
        let (s, e) = intervals[0];
        assert_eq!(e - s, 4 * 512);
    }
}