//! Speaker diarization (Phase 4, FR-SPK-*). sherpa-onnx: segmentation +
//! embedding + clustering, fully offline (ADR-0005).
//!
//! Runs as a post-stop pass over the recorded WAV (audio = source of truth),
//! then aligns speaker spans to transcript segments by timestamp overlap.
//! Speaker IDs are internal ("S1"…); names are applied at render time and never
//! destructively rewritten onto segments (FR-SPK-5).

use crate::models::{SpeakerSpan, TranscriptSegment};
use std::path::Path;

pub mod models;
pub mod voiceprint;

#[derive(Debug, thiserror::Error)]
pub enum DiarError {
    #[error("model load failed: {0}")]
    Load(String),
    #[error("diarization failed: {0}")]
    Run(String),
}

pub trait Diarizer: Send + Sync {
    /// Partition audio into speaker spans.
    fn diarize(&self, wav: &Path) -> Result<Vec<SpeakerSpan>, DiarError>;
    /// Assign speaker labels to transcript segments by overlap with spans.
    fn assign(&self, segments: &mut [TranscriptSegment], spans: &[SpeakerSpan]);
}

/// Labels each segment with whichever span overlaps it most, in milliseconds
/// (T4.2, FR-SPK-1). A segment with no overlapping span (e.g. it falls in a
/// gap between spans) keeps its prior speaker label — the "S1" placeholder
/// every segment starts with pre-diarization — rather than guessing.
///
/// Pure timestamp arithmetic, so it's engine-agnostic: every `Diarizer` impl
/// can share it instead of reimplementing overlap math.
pub fn assign_by_overlap(segments: &mut [TranscriptSegment], spans: &[SpeakerSpan]) {
    for segment in segments.iter_mut() {
        let best_span = spans
            .iter()
            .map(|span| {
                let overlap_start = segment.start_ms.max(span.start_ms);
                let overlap_end = segment.end_ms.min(span.end_ms);
                (overlap_end.saturating_sub(overlap_start), span)
            })
            .filter(|(overlap, _)| *overlap > 0)
            .max_by_key(|(overlap, _)| *overlap)
            .map(|(_, span)| span);

        if let Some(span) = best_span {
            segment.speaker = span.speaker.clone();
        }
    }
}

/// sherpa-onnx-backed diarizer: pyannote segmentation + speaker-embedding +
/// fast clustering (ADR-0005, T4.1). `Diarize::compute` needs `&mut self`; it's
/// wrapped in a `Mutex` to satisfy `Diarizer: Sync` — diarization is a
/// once-per-meeting post-pass (never a hot path), so lock contention is moot.
#[cfg(feature = "diarization")]
pub struct SherpaDiarizer {
    engine: std::sync::Mutex<sherpa_rs::diarize::Diarize>,
}

#[cfg(feature = "diarization")]
impl SherpaDiarizer {
    pub fn new(segmentation_model: &Path, embedding_model: &Path) -> Result<Self, DiarError> {
        let config = sherpa_rs::diarize::DiarizeConfig {
            // A meeting's speaker count isn't known ahead of time: <= 0 tells
            // sherpa-onnx to pick the cluster count itself from `threshold`
            // instead of forcing a fixed number of speakers.
            num_clusters: Some(-1),
            // Tuning evidence: 0.5 over the summed mic+loopback mono WAV
            // exploded a 2-speaker conversation into 83 clusters (overlapped
            // speech + short pyannote chunks yield unstable embeddings). sherpa
            // fast-clustering: larger threshold → fewer clusters. 0.7 lands a
            // 2-person meeting near 2–4. ponytail: code constant, no settings
            // knob until 0.7 proves insufficient in the field.
            threshold: Some(0.7),
            ..Default::default()
        };
        let engine = sherpa_rs::diarize::Diarize::new(segmentation_model, embedding_model, config)
            .map_err(|e| DiarError::Load(e.to_string()))?;
        Ok(Self {
            engine: std::sync::Mutex::new(engine),
        })
    }
}

#[cfg(feature = "diarization")]
impl Diarizer for SherpaDiarizer {
    fn diarize(&self, wav: &Path) -> Result<Vec<SpeakerSpan>, DiarError> {
        let samples =
            crate::audio::read_wav_mono_16k(wav).map_err(|e| DiarError::Run(e.to_string()))?;
        let mut engine = self.engine.lock().unwrap_or_else(|e| e.into_inner());
        let segments = engine
            .compute(samples, None)
            .map_err(|e| DiarError::Run(e.to_string()))?;
        // Drop sub-700ms spans: their embeddings are unstable and only cause
        // label churn in assign_by_overlap (which keeps the prior label when no
        // span overlaps). See MIN_SPAN_MS.
        Ok(segments
            .into_iter()
            .map(segment_to_span)
            .filter(span_long_enough)
            .collect())
    }

    fn assign(&self, segments: &mut [TranscriptSegment], spans: &[SpeakerSpan]) {
        assign_by_overlap(segments, spans);
    }
}

/// Minimum span duration kept from a diarization pass. Sub-second pyannote
/// chunks carry unstable ERes2Net embeddings and only churn labels.
#[cfg(feature = "diarization")]
const MIN_SPAN_MS: u64 = 700;

/// A span is kept only if it lasts at least [`MIN_SPAN_MS`].
#[cfg(feature = "diarization")]
fn span_long_enough(span: &SpeakerSpan) -> bool {
    span.end_ms.saturating_sub(span.start_ms) >= MIN_SPAN_MS
}

/// sherpa-onnx speaker indices are 0-based; WA's internal labels are 1-based ("S1"…).
#[cfg(feature = "diarization")]
fn segment_to_span(seg: sherpa_rs::diarize::Segment) -> SpeakerSpan {
    SpeakerSpan {
        start_ms: (seg.start.max(0.0) * 1000.0) as u64,
        end_ms: (seg.end.max(0.0) * 1000.0) as u64,
        speaker: format!("S{}", seg.speaker + 1),
    }
}

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

    fn segment(start_ms: u64, end_ms: u64) -> TranscriptSegment {
        TranscriptSegment {
            id: 0,
            start_ms,
            end_ms,
            speaker: "S1".to_string(),
            text: String::new(),
            confidence: None,
            interim: false,
        }
    }

    fn span(start_ms: u64, end_ms: u64, speaker: &str) -> SpeakerSpan {
        SpeakerSpan {
            start_ms,
            end_ms,
            speaker: speaker.to_string(),
        }
    }

    #[test]
    fn assigns_the_span_that_overlaps_most() {
        let mut segments = vec![segment(0, 1000), segment(1000, 2000)];
        let spans = vec![span(0, 1000, "S1"), span(1000, 2000, "S2")];
        assign_by_overlap(&mut segments, &spans);
        assert_eq!(segments[0].speaker, "S1");
        assert_eq!(segments[1].speaker, "S2");
    }

    #[test]
    fn a_segment_straddling_two_spans_picks_the_larger_overlap() {
        // 700ms in S1's span (300-1000), 300ms in S2's span (1000-1300).
        let mut segments = vec![segment(300, 1300)];
        let spans = vec![span(0, 1000, "S1"), span(1000, 2000, "S2")];
        assign_by_overlap(&mut segments, &spans);
        assert_eq!(segments[0].speaker, "S1");
    }

    #[test]
    fn a_segment_with_no_overlapping_span_keeps_its_prior_label() {
        let mut segments = vec![segment(5000, 6000)];
        segments[0].speaker = "S9".to_string(); // distinct from any span below
        let spans = vec![span(0, 1000, "S1")];
        assign_by_overlap(&mut segments, &spans);
        assert_eq!(segments[0].speaker, "S9"); // untouched, not overwritten with a guess
    }

    #[test]
    fn no_spans_at_all_leaves_segments_untouched() {
        let mut segments = vec![segment(0, 1000)];
        assign_by_overlap(&mut segments, &[]);
        assert_eq!(segments[0].speaker, "S1");
    }
}

#[cfg(all(test, feature = "diarization"))]
mod tests {
    use super::*;

    #[test]
    fn segment_to_span_maps_0_based_speaker_index_to_1_based_label() {
        let seg = sherpa_rs::diarize::Segment {
            start: 1.5,
            end: 3.25,
            speaker: 0,
        };
        let span = segment_to_span(seg);
        assert_eq!(span.start_ms, 1500);
        assert_eq!(span.end_ms, 3250);
        assert_eq!(span.speaker, "S1");
    }

    #[test]
    fn span_long_enough_filters_sub_minimum_spans() {
        let short = SpeakerSpan {
            start_ms: 0,
            end_ms: MIN_SPAN_MS - 1,
            speaker: "S1".to_string(),
        };
        let exact = SpeakerSpan {
            start_ms: 0,
            end_ms: MIN_SPAN_MS,
            speaker: "S1".to_string(),
        };
        assert!(!span_long_enough(&short));
        assert!(span_long_enough(&exact));
    }

    #[test]
    fn new_surfaces_a_load_error_for_missing_models_instead_of_panicking() {
        let result =
            SherpaDiarizer::new(Path::new("no-such-seg.onnx"), Path::new("no-such-emb.onnx"));
        assert!(matches!(result, Err(DiarError::Load(_))));
    }
}