car-voice 0.15.1

//! Speaker enrollment + per-segment role classification.
//!
//! A "voiceprint" is a fixed-dimension embedding of the user's voice
//! derived from a short enrollment utterance. At runtime the
//! [`SpeakerPipeline`] compares each fresh speech segment's embedding
//! against the stored voiceprint; close matches are tagged
//! [`TranscriptRole::EnrolledUser`], far matches [`TranscriptRole::OtherSpeaker`],
//! anything else [`TranscriptRole::Unknown`].
//!
//! The baseline embedder is mel-filterbank statistics (40 filters → 80-dim
//! mean+stddev), a pure-Rust zero-dep DSP good enough to discriminate
//! distinctly different speakers. Swap in ECAPA-TDNN via [`SpeakerPipeline::with_embedder`]
//! when an ONNX backend is wired up; the rest of the pipeline stays
//! the same.
//!
//! ## Wire format
//!
//! `Enrollment` persists as TOML (extension `.toml` by convention).
//! Fields use `serde(default)` where it keeps older files readable.
//!
//! ## Sample-rate invariant
//!
//! The trait signature [`SpeakerEmbedder::embed`] takes `&[i16]` alone
//! — the listener only ever feeds 16 kHz PCM (tokhn's codec constant)
//! so the DSP hardcodes that rate. If a future caller wants a
//! different rate, add a rate-aware variant; do not silently
//! reinterpret the existing one.

use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};

use crate::error::VoiceError;

/// Listeners convert f32 → i16 before embedding. This is the assumed
/// rate of the resulting PCM. Changing it silently would invalidate
/// every persisted voiceprint, so it's a named constant.
const EMBED_SAMPLE_RATE: u32 = 16_000;

/// Which speaker a transcript came from, as judged by the pipeline.
///
/// Default is [`TranscriptRole::Unknown`] — either because no pipeline
/// was configured or because the embedding distance fell between the
/// match and reject thresholds.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum TranscriptRole {
    /// No speaker attribution — pipeline disabled, below match
    /// threshold, or the segment was too short to embed.
    Unknown,
    /// Matched the enrolled voiceprint within `match_threshold`.
    EnrolledUser,
    /// Confidently identified as *not* the enrolled user. `local_id`
    /// is a per-session cluster handle so the consumer can thread
    /// "was this the same non-enrolled voice as last time?" without
    /// persisting anything. The one-shot classifier returns
    /// `"overheard"` — real per-turn diarization yields
    /// `"speaker_0"` / `"speaker_1"` / …
    OtherSpeaker { local_id: String },
}

impl Default for TranscriptRole {
    fn default() -> Self {
        TranscriptRole::Unknown
    }
}

/// Trait implemented by any component that turns 16 kHz mono PCM into
/// a fixed-length embedding vector. The vector space is implementation-
/// defined; the only contract is that two embeddings from the same
/// embedder are comparable via cosine similarity.
pub trait SpeakerEmbedder: Send + Sync {
    /// Returns the embedding for a 16-bit signed PCM segment at
    /// 16 kHz. An empty vector means the embedder could not produce
    /// a result (too short, too quiet, unsupported sample layout).
    fn embed(&self, pcm: &[i16]) -> Vec<f32>;

    /// Human-readable model identifier, stored alongside the embedding
    /// so loaders can reject voiceprints from incompatible models.
    fn model_id(&self) -> &str;
}

/// Filterbank-based speaker embedder: 40 mel-filter energies computed
/// per 25 ms frame with 10 ms hop, summarized as (mean, stddev) across
/// frames → 80-dim unit vector.
///
/// Pure Rust, zero model download, reasonable discrimination between
/// distinctly different speakers. Not production-grade — swap in an
/// ECAPA-TDNN ONNX backend via [`SpeakerPipeline::with_embedder`]
/// when the model bundling story is ready.
pub struct FilterbankEmbedder {
    model_id: String,
    mel_filters: usize,
}

impl FilterbankEmbedder {
    pub fn new() -> Self {
        Self {
            model_id: "fbank-stats-v1".to_string(),
            mel_filters: 40,
        }
    }
}

impl Default for FilterbankEmbedder {
    fn default() -> Self {
        Self::new()
    }
}

impl SpeakerEmbedder for FilterbankEmbedder {
    fn embed(&self, pcm: &[i16]) -> Vec<f32> {
        if pcm.is_empty() {
            return Vec::new();
        }
        let samples: Vec<f32> = pcm.iter().map(|&s| s as f32 / 32768.0).collect();
        let frames = mel_filterbank_frames(&samples, EMBED_SAMPLE_RATE, self.mel_filters);
        if frames.is_empty() {
            return Vec::new();
        }
        let mut means = vec![0.0f32; self.mel_filters];
        for frame in &frames {
            for (i, v) in frame.iter().enumerate() {
                means[i] += v;
            }
        }
        let n = frames.len() as f32;
        for m in &mut means {
            *m /= n;
        }
        let mut vars = vec![0.0f32; self.mel_filters];
        for frame in &frames {
            for (i, v) in frame.iter().enumerate() {
                let d = v - means[i];
                vars[i] += d * d;
            }
        }
        for v in &mut vars {
            *v = (*v / n).sqrt();
        }
        let mut values: Vec<f32> = Vec::with_capacity(self.mel_filters * 2);
        values.extend_from_slice(&means);
        values.extend_from_slice(&vars);
        normalize_in_place(&mut values);
        values
    }

    fn model_id(&self) -> &str {
        &self.model_id
    }
}

/// Voiceprint payload inside an [`Enrollment`]. Split into its own
/// struct so callers can inspect dimension (`values.len()`) and the
/// producing model without unpacking the whole enrollment.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpeakerEmbedding {
    #[serde(default)]
    pub values: Vec<f32>,
    #[serde(default)]
    pub model: String,
}

impl SpeakerEmbedding {
    pub fn new(values: Vec<f32>, model: impl Into<String>) -> Self {
        Self {
            values,
            model: model.into(),
        }
    }

    /// Length-normalize in place so cosine similarity reduces to dot
    /// product. [`FilterbankEmbedder`] already returns unit vectors;
    /// provided for callers that build embeddings from other sources.
    pub fn normalize(&mut self) {
        normalize_in_place(&mut self.values);
    }

    /// Cosine similarity between two embeddings. Refuses to compare
    /// across model families (different `model` strings) — cross-model
    /// distances are meaningless.
    pub fn cosine_similarity(
        a: &SpeakerEmbedding,
        b: &SpeakerEmbedding,
    ) -> Result<f32, VoiceError> {
        if a.model != b.model {
            return Err(VoiceError::Config(format!(
                "speaker embedding model mismatch: {} vs {}",
                a.model, b.model
            )));
        }
        if a.values.len() != b.values.len() {
            return Err(VoiceError::Config(format!(
                "speaker embedding dim mismatch: {} vs {}",
                a.values.len(),
                b.values.len()
            )));
        }
        if a.values.is_empty() {
            return Ok(0.0);
        }
        let dot: f32 = a
            .values
            .iter()
            .zip(b.values.iter())
            .map(|(x, y)| x * y)
            .sum();
        let na: f32 = a.values.iter().map(|v| v * v).sum::<f32>().sqrt();
        let nb: f32 = b.values.iter().map(|v| v * v).sum::<f32>().sqrt();
        let denom = na * nb;
        if denom <= f32::EPSILON {
            return Ok(0.0);
        }
        Ok(dot / denom)
    }
}

impl Default for SpeakerEmbedding {
    fn default() -> Self {
        Self {
            values: Vec::new(),
            model: String::new(),
        }
    }
}

/// A saved speaker profile. `label` is the user-facing name; the
/// embedding is compared to fresh segments at classify time.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Enrollment {
    pub label: String,
    #[serde(default)]
    pub embedding: SpeakerEmbedding,
}

impl Enrollment {
    /// Read a TOML-encoded enrollment from disk.
    pub fn load_from(path: &Path) -> Result<Self, VoiceError> {
        let bytes = std::fs::read(path)?;
        let text = std::str::from_utf8(&bytes)
            .map_err(|e| VoiceError::Config(format!("enrollment not UTF-8: {e}")))?;
        toml::from_str(text).map_err(|e| VoiceError::Config(format!("enrollment parse: {e}")))
    }

    /// Write this enrollment as TOML. Creates parent directories if
    /// they don't exist so the caller doesn't have to coordinate.
    pub fn save_to(&self, path: &Path) -> Result<(), VoiceError> {
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let text = toml::to_string(self)
            .map_err(|e| VoiceError::Config(format!("enrollment serialize: {e}")))?;
        std::fs::write(path, text)?;
        Ok(())
    }
}

/// An in-flight enrollment request. Armed via
/// [`SpeakerPipeline::arm_enrollment`]; the pipeline consumes it the
/// next time it sees enough audio to embed.
#[derive(Debug, Clone)]
pub struct PendingEnrollment {
    pub label: String,
    pub save_path: PathBuf,
    pub threshold: f32,
}

impl PendingEnrollment {
    pub fn new(label: impl Into<String>, save_path: impl Into<PathBuf>, threshold: f32) -> Self {
        Self {
            label: label.into(),
            save_path: save_path.into(),
            threshold,
        }
    }
}

/// Outcome of running a speech segment through an armed enrollment
/// capture. Listeners translate this into the matching [`crate::VoiceEvent`]
/// (captured vs failed).
#[derive(Debug, Clone)]
pub enum EnrollmentOutcome {
    Captured { label: String, save_path: PathBuf },
    Failed { reason: String },
}

/// Default match-threshold cosine similarity when none is provided.
/// Values above this → enrolled user. Values well below → other
/// speaker. Tuned for [`FilterbankEmbedder`]; ECAPA-TDNN typically
/// runs closer to 0.45.
const DEFAULT_MATCH_THRESHOLD: f32 = 0.72;

/// Cosine similarity gap below the match threshold at which we stop
/// calling a segment "other speaker" and fall back to `Unknown`.
/// Guards against short/quiet segments whose embedding noise happens
/// to land well below threshold for the enrolled user too.
const OTHER_SPEAKER_MARGIN: f32 = 0.15;

// ─────────────────────────────────────────────────────────────────────────────
// Diarization types + traits
// ─────────────────────────────────────────────────────────────────────────────

/// One contiguous, single-speaker span within a diarized segment.
///
/// `speaker_id` is local to the diarization call — it is *not* stable
/// across calls. Use it only to group turns within the same segment.
/// Cross-segment identity comes from comparing `embedding` to an
/// [`Enrollment`].
#[derive(Debug, Clone)]
pub struct DiarizationTurn {
    pub start_ms: u64,
    pub end_ms: u64,
    pub speaker_id: String,
    pub embedding: SpeakerEmbedding,
}

/// Who spoke a given turn. The per-turn analogue of [`TranscriptRole`];
/// preserves the raw diarizer cluster id so downstream consumers can
/// group consecutive turns from the same non-enrolled speaker.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TurnRole {
    EnrolledUser,
    OtherSpeaker(String),
    Unknown,
}

impl TurnRole {
    /// Convert to the segment-level [`TranscriptRole`] emitted on
    /// [`crate::VoiceEvent::Transcript`]. The `OtherSpeaker` cluster
    /// id rides through unchanged.
    pub fn into_transcript_role(self) -> TranscriptRole {
        match self {
            TurnRole::EnrolledUser => TranscriptRole::EnrolledUser,
            TurnRole::OtherSpeaker(id) => TranscriptRole::OtherSpeaker { local_id: id },
            TurnRole::Unknown => TranscriptRole::Unknown,
        }
    }
}

/// A diarization turn tagged with its [`TurnRole`] plus the per-turn
/// PCM. Returned by [`SpeakerPipeline::classify_turns`] so callers can
/// feed each turn straight to STT without re-slicing the original
/// segment.
#[derive(Debug, Clone)]
pub struct TaggedTurn {
    pub start_ms: u64,
    pub end_ms: u64,
    pub role: TurnRole,
    pub embedding: SpeakerEmbedding,
    /// Normalized f32 samples in `[-1.0, 1.0]` — the format STT
    /// providers consume directly.
    pub audio: Vec<f32>,
}

/// Decide whether audio matches an enrollment.
pub trait SpeakerVerifier: Send + Sync {
    /// Score-then-threshold check. `threshold` comes from the caller
    /// (usually [`SpeakerPipeline::match_threshold`]) rather than the
    /// enrollment so the pipeline can vary strictness per call site.
    fn matches(
        &self,
        samples: &[f32],
        sample_rate: u32,
        enrolled: &Enrollment,
        threshold: f32,
    ) -> Result<bool, VoiceError>;

    /// Raw cosine similarity without the threshold check. Useful for
    /// logging + threshold tuning.
    fn score(
        &self,
        samples: &[f32],
        sample_rate: u32,
        enrolled: &Enrollment,
    ) -> Result<f32, VoiceError>;
}

/// Split mixed-speaker audio into per-speaker turns.
pub trait Diarizer: Send + Sync {
    fn diarize(
        &self,
        samples: &[f32],
        sample_rate: u32,
    ) -> Result<Vec<DiarizationTurn>, VoiceError>;
}

/// Verifier that embeds the sample, cosine-compares to the enrolled
/// embedding, and returns `score >= threshold`. Works with any
/// [`SpeakerEmbedder`] as long as the model labels match.
pub struct CosineVerifier {
    embedder: Arc<dyn SpeakerEmbedder>,
}

impl CosineVerifier {
    pub fn new(embedder: Arc<dyn SpeakerEmbedder>) -> Self {
        Self { embedder }
    }
}

impl SpeakerVerifier for CosineVerifier {
    fn matches(
        &self,
        samples: &[f32],
        sample_rate: u32,
        enrolled: &Enrollment,
        threshold: f32,
    ) -> Result<bool, VoiceError> {
        Ok(self.score(samples, sample_rate, enrolled)? >= threshold)
    }

    fn score(
        &self,
        samples: &[f32],
        sample_rate: u32,
        enrolled: &Enrollment,
    ) -> Result<f32, VoiceError> {
        if sample_rate != EMBED_SAMPLE_RATE {
            return Err(VoiceError::Config(format!(
                "CosineVerifier expects {EMBED_SAMPLE_RATE} Hz; got {sample_rate}"
            )));
        }
        let pcm = f32_to_i16(samples);
        let values = self.embedder.embed(&pcm);
        if values.is_empty() {
            return Err(VoiceError::Config("embedder produced empty vector".into()));
        }
        let candidate = SpeakerEmbedding::new(values, self.embedder.model_id());
        SpeakerEmbedding::cosine_similarity(&candidate, &enrolled.embedding)
    }
}

/// Silence-split + agglomerative-merge diarizer. Splits the segment
/// at silence boundaries, embeds each voiced chunk, and clusters by
/// cosine similarity ≥ `merge_threshold`.
///
/// Good enough for "two alternating speakers in a short segment";
/// pathological cases (overlapping voices mid-frame) need a real
/// ONNX model. Returns up to N turns where N is the number of
/// distinct voiced runs in the segment.
pub struct SilenceSplitDiarizer {
    embedder: Arc<dyn SpeakerEmbedder>,
    merge_threshold: f32,
    silence_threshold_dbfs: f32,
    min_chunk_ms: u64,
}

impl SilenceSplitDiarizer {
    pub fn new(embedder: Arc<dyn SpeakerEmbedder>) -> Self {
        Self {
            embedder,
            merge_threshold: 0.75,
            silence_threshold_dbfs: -45.0,
            min_chunk_ms: 300,
        }
    }

    pub fn with_merge_threshold(mut self, t: f32) -> Self {
        self.merge_threshold = t;
        self
    }

    pub fn with_silence_threshold_dbfs(mut self, t: f32) -> Self {
        self.silence_threshold_dbfs = t;
        self
    }

    pub fn with_min_chunk_ms(mut self, ms: u64) -> Self {
        self.min_chunk_ms = ms;
        self
    }
}

impl Diarizer for SilenceSplitDiarizer {
    fn diarize(
        &self,
        samples: &[f32],
        sample_rate: u32,
    ) -> Result<Vec<DiarizationTurn>, VoiceError> {
        if sample_rate != EMBED_SAMPLE_RATE {
            return Err(VoiceError::Config(format!(
                "SilenceSplitDiarizer expects {EMBED_SAMPLE_RATE} Hz; got {sample_rate}"
            )));
        }
        let chunks = split_on_silence(
            samples,
            sample_rate,
            self.silence_threshold_dbfs,
            self.min_chunk_ms,
        );
        if chunks.is_empty() {
            return Ok(Vec::new());
        }
        let mut embedded: Vec<(u64, u64, SpeakerEmbedding)> = Vec::new();
        for (start_ms, end_ms, piece) in &chunks {
            let pcm = f32_to_i16(piece);
            let values = self.embedder.embed(&pcm);
            if values.is_empty() {
                continue;
            }
            let emb = SpeakerEmbedding::new(values, self.embedder.model_id());
            embedded.push((*start_ms, *end_ms, emb));
        }
        // Agglomerative: walk chunks, assign to the best existing
        // cluster if similar enough, else open a new one.
        let mut clusters: Vec<SpeakerEmbedding> = Vec::new();
        let mut assignments: Vec<usize> = Vec::with_capacity(embedded.len());
        for (_, _, emb) in &embedded {
            let mut best: Option<(usize, f32)> = None;
            for (i, c) in clusters.iter().enumerate() {
                if let Ok(s) = SpeakerEmbedding::cosine_similarity(c, emb) {
                    if best.map(|(_, bs)| s > bs).unwrap_or(true) {
                        best = Some((i, s));
                    }
                }
            }
            match best {
                Some((i, s)) if s >= self.merge_threshold => assignments.push(i),
                _ => {
                    clusters.push(emb.clone());
                    assignments.push(clusters.len() - 1);
                }
            }
        }
        Ok(embedded
            .into_iter()
            .zip(assignments.iter())
            .map(
                |((start_ms, end_ms, embedding), cluster_id)| DiarizationTurn {
                    start_ms,
                    end_ms,
                    speaker_id: format!("speaker_{cluster_id}"),
                    embedding,
                },
            )
            .collect())
    }
}

/// Composes an embedder, an optional loaded [`Enrollment`], and a
/// possibly-armed [`PendingEnrollment`] into a single handle the
/// listener can consult per segment.
///
/// **Shared across threads.** The listener's capture thread calls
/// `classify`/`capture_enrollment` while the bridge thread may call
/// `arm_enrollment` — hence the interior `Mutex`. Uncontended in
/// practice (one reader, one writer, no overlap).
pub struct SpeakerPipeline {
    embedder: Arc<dyn SpeakerEmbedder>,
    verifier: Arc<dyn SpeakerVerifier>,
    diarizer: Arc<dyn Diarizer>,
    enrollment: Mutex<Option<Enrollment>>,
    match_threshold: f32,
    pending: Mutex<Option<PendingEnrollment>>,
}

impl SpeakerPipeline {
    /// Pipeline with the filterbank embedder, a cosine verifier, a
    /// silence-split diarizer, no enrollment, and the default match
    /// threshold.
    pub fn baseline() -> Self {
        let embedder: Arc<dyn SpeakerEmbedder> = Arc::new(FilterbankEmbedder::new());
        let verifier: Arc<dyn SpeakerVerifier> =
            Arc::new(CosineVerifier::new(Arc::clone(&embedder)));
        let diarizer: Arc<dyn Diarizer> =
            Arc::new(SilenceSplitDiarizer::new(Arc::clone(&embedder)));
        Self {
            embedder,
            verifier,
            diarizer,
            enrollment: Mutex::new(None),
            match_threshold: DEFAULT_MATCH_THRESHOLD,
            pending: Mutex::new(None),
        }
    }

    /// Replace the embedder. Useful once an ECAPA-TDNN (or similar)
    /// backend lands — the rest of the pipeline stays the same.
    ///
    /// Note this replaces only the pipeline's classify-time embedder.
    /// If you also want the verifier + diarizer to use the new
    /// embedder, pass them explicitly via
    /// [`SpeakerPipeline::with_verifier`] / [`SpeakerPipeline::with_diarizer`].
    pub fn with_embedder(mut self, embedder: Box<dyn SpeakerEmbedder>) -> Self {
        self.embedder = Arc::from(embedder);
        self
    }

    /// Swap in a custom verifier — useful for A/B-testing alternate
    /// scoring functions against the same enrollment.
    pub fn with_verifier(mut self, verifier: Arc<dyn SpeakerVerifier>) -> Self {
        self.verifier = verifier;
        self
    }

    /// Swap in a custom diarizer — e.g. an ONNX model once one is
    /// wired up.
    pub fn with_diarizer(mut self, diarizer: Arc<dyn Diarizer>) -> Self {
        self.diarizer = diarizer;
        self
    }

    /// Attach a loaded voiceprint so subsequent `classify` calls can
    /// compare against it.
    pub fn with_enrollment(self, enrollment: Enrollment) -> Self {
        if let Ok(mut slot) = self.enrollment.lock() {
            *slot = Some(enrollment);
        }
        self
    }

    /// Override the cosine-similarity threshold. Higher = stricter
    /// "is this the enrolled user".
    pub fn with_match_threshold(mut self, threshold: f32) -> Self {
        self.match_threshold = threshold;
        self
    }

    /// Snapshot of the currently-loaded enrollment. Cheap clone;
    /// used by `classify` and also useful for GUIs that render "who
    /// is enrolled" indicators.
    pub fn enrollment_snapshot(&self) -> Option<Enrollment> {
        self.enrollment.lock().ok().and_then(|s| s.clone())
    }

    /// Classify a fresh speech segment.
    ///
    /// Returns `Unknown` when no enrollment is configured, the
    /// embedder produced an empty vector, or the cosine similarity
    /// sits in the ambiguous band between `match_threshold` and
    /// `match_threshold - OTHER_SPEAKER_MARGIN`. Otherwise returns
    /// `EnrolledUser` (≥ threshold) or `OtherSpeaker { "overheard" }`
    /// (well below threshold).
    pub fn classify(&self, pcm: &[i16]) -> TranscriptRole {
        let Some(enrollment) = self.enrollment_snapshot() else {
            return TranscriptRole::Unknown;
        };
        if enrollment.embedding.values.is_empty() {
            return TranscriptRole::Unknown;
        }
        let values = self.embedder.embed(pcm);
        if values.is_empty() {
            return TranscriptRole::Unknown;
        }
        let candidate = SpeakerEmbedding::new(values, self.embedder.model_id());
        let score = match SpeakerEmbedding::cosine_similarity(&candidate, &enrollment.embedding) {
            Ok(s) => s,
            Err(e) => {
                tracing::warn!("[voice] cosine compare rejected: {e}");
                return TranscriptRole::Unknown;
            }
        };
        if score >= self.match_threshold {
            TranscriptRole::EnrolledUser
        } else if score < self.match_threshold - OTHER_SPEAKER_MARGIN {
            TranscriptRole::OtherSpeaker {
                local_id: "overheard".to_string(),
            }
        } else {
            TranscriptRole::Unknown
        }
    }

    /// Diarize a segment and classify every turn against the loaded
    /// enrollment. Returns the per-turn audio alongside each tag so
    /// the caller can feed each turn straight to STT.
    ///
    /// Returns an empty vector when no enrollment is configured — the
    /// caller should then fall back to treating the whole segment as
    /// one `Unknown` turn (the pre-enrollment path).
    pub fn classify_turns(
        &self,
        samples: &[f32],
        sample_rate: u32,
    ) -> Result<Vec<TaggedTurn>, VoiceError> {
        let Some(enrollment) = self.enrollment_snapshot() else {
            return Ok(Vec::new());
        };
        let turns = self.diarizer.diarize(samples, sample_rate)?;
        if turns.is_empty() {
            return Ok(Vec::new());
        }
        let mut out: Vec<TaggedTurn> = Vec::with_capacity(turns.len());
        for turn in turns {
            let role =
                match SpeakerEmbedding::cosine_similarity(&turn.embedding, &enrollment.embedding) {
                    Ok(score) if score >= self.match_threshold => TurnRole::EnrolledUser,
                    Ok(_) => TurnRole::OtherSpeaker(turn.speaker_id.clone()),
                    Err(e) => {
                        tracing::warn!("[voice] cosine compare rejected in classify_turns: {e}");
                        TurnRole::Unknown
                    }
                };
            let start = ((turn.start_ms * sample_rate as u64) / 1000) as usize;
            let end = ((turn.end_ms * sample_rate as u64) / 1000) as usize;
            let end = end.min(samples.len());
            let audio = if end > start {
                samples[start..end].to_vec()
            } else {
                Vec::new()
            };
            out.push(TaggedTurn {
                start_ms: turn.start_ms,
                end_ms: turn.end_ms,
                role,
                embedding: turn.embedding,
                audio,
            });
        }
        Ok(out)
    }

    /// Drop non-enrolled speakers from an audio segment.
    ///
    /// Returns `Some(audio)` containing only the enrolled user's
    /// speech (concatenated with short silence pads between kept
    /// turns), `None` if no enrolled-user turn was found — the caller
    /// should then drop the whole segment — or `Some(original)` when
    /// no enrollment is configured (pass-through).
    pub fn filter_to_enrolled_user(
        &self,
        samples: &[f32],
        sample_rate: u32,
    ) -> Result<Option<Vec<f32>>, VoiceError> {
        let Some(enrollment) = self.enrollment_snapshot() else {
            return Ok(Some(samples.to_vec()));
        };

        // Fast path: one-speaker check against the full segment.
        if let Ok(score) = self.verifier.score(samples, sample_rate, &enrollment) {
            if score >= self.match_threshold {
                return Ok(Some(samples.to_vec()));
            }
        }

        let turns = self.diarizer.diarize(samples, sample_rate)?;
        if turns.is_empty() {
            return Ok(None);
        }
        let mut kept: Vec<Vec<f32>> = Vec::new();
        for turn in turns {
            let score = SpeakerEmbedding::cosine_similarity(&turn.embedding, &enrollment.embedding)
                .unwrap_or(0.0);
            if score >= self.match_threshold {
                let start = ((turn.start_ms * sample_rate as u64) / 1000) as usize;
                let end = ((turn.end_ms * sample_rate as u64) / 1000) as usize;
                let end = end.min(samples.len());
                if end > start {
                    kept.push(samples[start..end].to_vec());
                }
            }
        }
        if kept.is_empty() {
            return Ok(None);
        }
        // Concatenate retained turns with 100 ms of silence between
        // them so STT doesn't glue adjacent words from different
        // parts of the segment.
        let pad_samples = (sample_rate as usize) / 10;
        let pad = vec![0.0f32; pad_samples];
        let mut out: Vec<f32> = Vec::new();
        for (i, chunk) in kept.into_iter().enumerate() {
            if i > 0 {
                out.extend_from_slice(&pad);
            }
            out.extend(chunk);
        }
        Ok(Some(out))
    }

    /// Arm an enrollment capture. The next speech segment the
    /// listener feeds through [`SpeakerPipeline::capture_enrollment`]
    /// becomes the voiceprint.
    ///
    /// Replacing an existing pending request is fine — the newest arm
    /// wins.
    pub fn arm_enrollment(&self, request: PendingEnrollment) {
        if let Ok(mut guard) = self.pending.lock() {
            *guard = Some(request);
        }
    }

    /// Consume any armed enrollment, embed the provided PCM, write a
    /// voiceprint to the requested path, and hot-update the pipeline's
    /// loaded enrollment so subsequent `classify` calls attribute
    /// speech to the newly-enrolled user.
    ///
    /// Returns `None` when no enrollment is armed (listener should
    /// route the segment to STT normally).
    pub fn capture_enrollment(&self, pcm: &[i16]) -> Option<EnrollmentOutcome> {
        let pending = self.pending.lock().ok().and_then(|mut g| g.take())?;
        let values = self.embedder.embed(pcm);
        if values.is_empty() {
            return Some(EnrollmentOutcome::Failed {
                reason: format!(
                    "embedder '{}' returned empty vector — audio too short or silent",
                    self.embedder.model_id()
                ),
            });
        }
        let enrollment = Enrollment {
            label: pending.label.clone(),
            embedding: SpeakerEmbedding {
                values,
                model: self.embedder.model_id().to_string(),
            },
        };
        if let Err(e) = enrollment.save_to(&pending.save_path) {
            return Some(EnrollmentOutcome::Failed {
                reason: format!("save to {}: {e}", pending.save_path.display()),
            });
        }
        // Hot-update so the very next segment is classified against
        // the new voiceprint without requiring the caller to rebuild
        // the pipeline.
        if let Ok(mut slot) = self.enrollment.lock() {
            *slot = Some(enrollment);
        }
        Some(EnrollmentOutcome::Captured {
            label: pending.label,
            save_path: pending.save_path,
        })
    }

    /// Non-destructive peek so a listener can decide "do I need to
    /// route this segment through enrollment capture?" without
    /// committing to consuming the pending request.
    pub fn has_pending_enrollment(&self) -> bool {
        self.pending.lock().map(|g| g.is_some()).unwrap_or(false)
    }

    /// Whether the pipeline has an enrolled voiceprint to compare
    /// against. `classify` falls back to `Unknown` when this is
    /// `false` regardless of input.
    pub fn is_enrolled(&self) -> bool {
        self.enrollment.lock().map(|s| s.is_some()).unwrap_or(false)
    }
}

impl Default for SpeakerPipeline {
    fn default() -> Self {
        Self::baseline()
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// DSP helpers
// ─────────────────────────────────────────────────────────────────────────────

fn normalize_in_place(values: &mut [f32]) {
    let norm: f32 = values.iter().map(|v| v * v).sum::<f32>().sqrt();
    if norm > f32::EPSILON {
        for v in values {
            *v /= norm;
        }
    }
}

/// Convert normalized `[-1.0, 1.0]` f32 samples to 16-bit signed PCM
/// — the format [`SpeakerEmbedder::embed`] takes. Clamps out-of-range
/// values (rare; AGC keeps things below unity).
fn f32_to_i16(samples: &[f32]) -> Vec<i16> {
    samples
        .iter()
        .map(|&s| (s.clamp(-1.0, 1.0) * 32767.0).round() as i16)
        .collect()
}

/// Silence-based splitter. Returns `(start_ms, end_ms, samples)` per
/// voiced run of at least `min_chunk_ms`. Frame size is 20 ms; RMS
/// below `threshold_dbfs` counts as silence.
fn split_on_silence(
    samples: &[f32],
    sample_rate: u32,
    threshold_dbfs: f32,
    min_chunk_ms: u64,
) -> Vec<(u64, u64, Vec<f32>)> {
    if samples.is_empty() {
        return Vec::new();
    }
    let frame_size = ((sample_rate as f32) * 0.020).round() as usize;
    if frame_size == 0 {
        return Vec::new();
    }
    let threshold_amp = 10f32.powf(threshold_dbfs / 20.0);
    let mut out: Vec<(u64, u64, Vec<f32>)> = Vec::new();
    let mut chunk_start: Option<usize> = None;
    let mut i = 0usize;
    while i + frame_size <= samples.len() {
        let frame = &samples[i..i + frame_size];
        let rms = (frame.iter().map(|s| s * s).sum::<f32>() / frame.len() as f32).sqrt();
        let voiced = rms >= threshold_amp;
        match (voiced, chunk_start) {
            (true, None) => chunk_start = Some(i),
            (false, Some(start)) => {
                let ms = ((i - start) as u64 * 1000) / sample_rate as u64;
                if ms >= min_chunk_ms {
                    out.push((
                        (start as u64 * 1000) / sample_rate as u64,
                        (i as u64 * 1000) / sample_rate as u64,
                        samples[start..i].to_vec(),
                    ));
                }
                chunk_start = None;
            }
            _ => {}
        }
        i += frame_size;
    }
    if let Some(start) = chunk_start {
        let ms = ((samples.len() - start) as u64 * 1000) / sample_rate as u64;
        if ms >= min_chunk_ms {
            out.push((
                (start as u64 * 1000) / sample_rate as u64,
                (samples.len() as u64 * 1000) / sample_rate as u64,
                samples[start..].to_vec(),
            ));
        }
    }
    out
}

/// Compute mel-filterbank energies frame-by-frame. Short-time Fourier
/// transform windows @ 25 ms with 10 ms hop; `n_mels` triangular
/// filters across 0–8 kHz.
fn mel_filterbank_frames(samples: &[f32], sample_rate: u32, n_mels: usize) -> Vec<Vec<f32>> {
    let frame_size = ((sample_rate as f32) * 0.025).round() as usize;
    let hop_size = ((sample_rate as f32) * 0.010).round() as usize;
    if frame_size == 0 || hop_size == 0 || samples.len() < frame_size {
        return vec![];
    }
    let window: Vec<f32> = (0..frame_size)
        .map(|i| {
            0.5 - 0.5 * ((2.0 * std::f32::consts::PI * i as f32) / (frame_size - 1) as f32).cos()
        })
        .collect();

    let mel_min = hz_to_mel(80.0);
    let mel_max = hz_to_mel((sample_rate as f32 / 2.0).min(8000.0));
    let mel_edges: Vec<f32> = (0..n_mels + 2)
        .map(|i| mel_min + (mel_max - mel_min) * (i as f32) / (n_mels + 1) as f32)
        .collect();
    let hz_edges: Vec<f32> = mel_edges.iter().copied().map(mel_to_hz).collect();
    let n_fft = frame_size.next_power_of_two();
    let bin_edges: Vec<f32> = hz_edges
        .iter()
        .map(|hz| hz * (n_fft as f32) / (sample_rate as f32))
        .collect();

    let mut out: Vec<Vec<f32>> = Vec::new();
    let mut frame_samples = vec![0.0f32; n_fft];
    let mut start = 0usize;
    while start + frame_size <= samples.len() {
        for i in 0..n_fft {
            frame_samples[i] = if i < frame_size {
                samples[start + i] * window[i]
            } else {
                0.0
            };
        }
        let mags = dft_magnitudes(&frame_samples);
        let mut mels = vec![0.0f32; n_mels];
        for m in 0..n_mels {
            let lo = bin_edges[m];
            let ctr = bin_edges[m + 1];
            let hi = bin_edges[m + 2];
            let mut acc = 0.0f32;
            for (k, mag) in mags.iter().enumerate() {
                let k = k as f32;
                let weight = if k <= lo || k >= hi {
                    0.0
                } else if k <= ctr {
                    (k - lo) / (ctr - lo).max(1e-6)
                } else {
                    (hi - k) / (hi - ctr).max(1e-6)
                };
                acc += weight * mag;
            }
            mels[m] = (acc + 1e-6).ln();
        }
        out.push(mels);
        start += hop_size;
    }
    out
}

fn hz_to_mel(hz: f32) -> f32 {
    2595.0 * (1.0 + hz / 700.0).log10()
}

fn mel_to_hz(mel: f32) -> f32 {
    700.0 * (10f32.powf(mel / 2595.0) - 1.0)
}

/// Naive DFT magnitude spectrum — O(N²) but fine for 25 ms × 16 kHz =
/// 400 samples → 512 FFT size frames, called once per hop. Keeps
/// `car-voice` dependency-free of `rustfft` in this baseline pass.
fn dft_magnitudes(frame: &[f32]) -> Vec<f32> {
    let n = frame.len();
    let half = n / 2 + 1;
    let mut out = vec![0.0f32; half];
    for k in 0..half {
        let mut re = 0.0f32;
        let mut im = 0.0f32;
        let factor = -2.0 * std::f32::consts::PI * (k as f32) / (n as f32);
        for (t, x) in frame.iter().enumerate() {
            let angle = factor * (t as f32);
            re += x * angle.cos();
            im += x * angle.sin();
        }
        out[k] = (re * re + im * im).sqrt();
    }
    out
}

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

    fn tone_i16(freq: f32, duration_s: f32, sample_rate: u32) -> Vec<i16> {
        let total = (duration_s * sample_rate as f32) as usize;
        (0..total)
            .map(|i| {
                let t = i as f32 / sample_rate as f32;
                let s = (2.0 * std::f32::consts::PI * freq * t).sin() * 0.4;
                (s * 32767.0) as i16
            })
            .collect()
    }

    #[test]
    fn default_role_is_unknown() {
        assert_eq!(TranscriptRole::default(), TranscriptRole::Unknown);
    }

    #[test]
    fn embedder_returns_unit_vector_for_tone() {
        let audio = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        let embedder = FilterbankEmbedder::new();
        let values = embedder.embed(&audio);
        assert_eq!(values.len(), 80);
        let norm: f32 = values.iter().map(|v| v * v).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-4, "embedding not unit-norm: {norm}");
    }

    #[test]
    fn embedder_returns_empty_for_empty_input() {
        let embedder = FilterbankEmbedder::new();
        assert!(embedder.embed(&[]).is_empty());
    }

    #[test]
    fn same_tone_self_matches() {
        let audio = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        let embedder = FilterbankEmbedder::new();
        let a = SpeakerEmbedding::new(embedder.embed(&audio), embedder.model_id());
        let b = SpeakerEmbedding::new(embedder.embed(&audio), embedder.model_id());
        let score = SpeakerEmbedding::cosine_similarity(&a, &b).unwrap();
        assert!(
            score > 0.99,
            "identical audio should self-match, got {score}"
        );
    }

    #[test]
    fn different_tones_score_lower_than_same_tone() {
        let same_a = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        let same_b = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        let different = tone_i16(880.0, 1.0, EMBED_SAMPLE_RATE);
        let embedder = FilterbankEmbedder::new();
        let ea = SpeakerEmbedding::new(embedder.embed(&same_a), embedder.model_id());
        let eb = SpeakerEmbedding::new(embedder.embed(&same_b), embedder.model_id());
        let ed = SpeakerEmbedding::new(embedder.embed(&different), embedder.model_id());
        let same_score = SpeakerEmbedding::cosine_similarity(&ea, &eb).unwrap();
        let diff_score = SpeakerEmbedding::cosine_similarity(&ea, &ed).unwrap();
        assert!(
            same_score > diff_score + 0.05,
            "same: {same_score}, different: {diff_score}"
        );
    }

    #[test]
    fn cosine_similarity_refuses_model_mismatch() {
        let a = SpeakerEmbedding::new(vec![1.0, 0.0], "model-a");
        let b = SpeakerEmbedding::new(vec![1.0, 0.0], "model-b");
        assert!(SpeakerEmbedding::cosine_similarity(&a, &b).is_err());
    }

    #[test]
    fn baseline_without_enrollment_is_unknown() {
        let p = SpeakerPipeline::baseline();
        let audio = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        assert_eq!(p.classify(&audio), TranscriptRole::Unknown);
    }

    #[test]
    fn classify_returns_enrolled_user_for_matching_audio() {
        let embedder = FilterbankEmbedder::new();
        let audio = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        let enrolled = Enrollment {
            label: "matt".into(),
            embedding: SpeakerEmbedding::new(embedder.embed(&audio), embedder.model_id()),
        };
        let pipeline = SpeakerPipeline::baseline().with_enrollment(enrolled);
        assert_eq!(pipeline.classify(&audio), TranscriptRole::EnrolledUser);
    }

    #[test]
    fn classify_returns_other_for_very_different_audio() {
        let embedder = FilterbankEmbedder::new();
        let enrolled_audio = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        let other_audio = tone_i16(2_000.0, 1.0, EMBED_SAMPLE_RATE);
        let enrolled = Enrollment {
            label: "matt".into(),
            embedding: SpeakerEmbedding::new(embedder.embed(&enrolled_audio), embedder.model_id()),
        };
        let pipeline = SpeakerPipeline::baseline().with_enrollment(enrolled);
        match pipeline.classify(&other_audio) {
            TranscriptRole::OtherSpeaker { .. } | TranscriptRole::Unknown => {}
            TranscriptRole::EnrolledUser => {
                panic!("2 kHz tone should not match a 220 Hz enrollment")
            }
        }
    }

    #[test]
    fn enrollment_roundtrip() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("voiceprint.toml");
        let e = Enrollment {
            label: "matt".into(),
            embedding: SpeakerEmbedding {
                values: vec![0.1, 0.2, 0.3],
                model: "fbank-stats-v1".into(),
            },
        };
        e.save_to(&path).unwrap();
        let loaded = Enrollment::load_from(&path).unwrap();
        assert_eq!(loaded.label, "matt");
        assert_eq!(loaded.embedding.values.len(), 3);
        assert_eq!(loaded.embedding.model, "fbank-stats-v1");
    }

    #[test]
    fn arm_then_capture_produces_voiceprint() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("voiceprint.toml");
        let pipeline = SpeakerPipeline::baseline();
        assert!(!pipeline.has_pending_enrollment());
        pipeline.arm_enrollment(PendingEnrollment::new("matt", &path, 0.72));
        assert!(pipeline.has_pending_enrollment());

        let audio = tone_i16(220.0, 1.0, EMBED_SAMPLE_RATE);
        let outcome = pipeline.capture_enrollment(&audio).unwrap();
        match outcome {
            EnrollmentOutcome::Captured { label, save_path } => {
                assert_eq!(label, "matt");
                assert_eq!(save_path, path);
            }
            EnrollmentOutcome::Failed { reason } => {
                panic!("expected Captured, got Failed: {reason}")
            }
        }
        assert!(!pipeline.has_pending_enrollment());
        assert!(pipeline.is_enrolled());
        // Post-capture, the same audio should classify as the enrolled user.
        assert_eq!(pipeline.classify(&audio), TranscriptRole::EnrolledUser);
    }

    #[test]
    fn capture_with_empty_audio_fails_cleanly() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("voiceprint.toml");
        let pipeline = SpeakerPipeline::baseline();
        pipeline.arm_enrollment(PendingEnrollment::new("matt", &path, 0.72));
        let outcome = pipeline.capture_enrollment(&[]).unwrap();
        assert!(matches!(outcome, EnrollmentOutcome::Failed { .. }));
        assert!(!pipeline.has_pending_enrollment());
    }

    // ─────────────────────────────────────────────────────────────────
    // Diarizer / per-turn tests
    // ─────────────────────────────────────────────────────────────────

    fn tone_f32(freq: f32, duration_s: f32, sample_rate: u32) -> Vec<f32> {
        let total = (duration_s * sample_rate as f32) as usize;
        (0..total)
            .map(|i| {
                let t = i as f32 / sample_rate as f32;
                (2.0 * std::f32::consts::PI * freq * t).sin() * 0.4
            })
            .collect()
    }

    fn silence_f32(duration_s: f32, sample_rate: u32) -> Vec<f32> {
        vec![0.0f32; (duration_s * sample_rate as f32) as usize]
    }

    #[test]
    fn turn_role_maps_to_transcript_role() {
        assert_eq!(
            TurnRole::EnrolledUser.into_transcript_role(),
            TranscriptRole::EnrolledUser
        );
        assert_eq!(
            TurnRole::OtherSpeaker("speaker_3".into()).into_transcript_role(),
            TranscriptRole::OtherSpeaker {
                local_id: "speaker_3".into()
            }
        );
        assert_eq!(
            TurnRole::Unknown.into_transcript_role(),
            TranscriptRole::Unknown
        );
    }

    #[test]
    fn diarizer_splits_two_tones_separated_by_silence() {
        let sr = EMBED_SAMPLE_RATE;
        let mut audio = tone_f32(220.0, 1.0, sr);
        audio.extend(silence_f32(0.5, sr));
        audio.extend(tone_f32(1_500.0, 1.0, sr));
        let embedder: Arc<dyn SpeakerEmbedder> = Arc::new(FilterbankEmbedder::new());
        let diarizer = SilenceSplitDiarizer::new(embedder);
        let turns = diarizer.diarize(&audio, sr).unwrap();
        assert_eq!(turns.len(), 2, "expected 2 turns, got {}", turns.len());
        assert_ne!(
            turns[0].speaker_id, turns[1].speaker_id,
            "different tones should land in different clusters"
        );
    }

    #[test]
    fn classify_turns_empty_without_enrollment() {
        let pipeline = SpeakerPipeline::baseline();
        let audio = tone_f32(220.0, 1.0, EMBED_SAMPLE_RATE);
        let turns = pipeline.classify_turns(&audio, EMBED_SAMPLE_RATE).unwrap();
        assert!(turns.is_empty());
    }

    #[test]
    fn classify_turns_tags_enrolled_and_other() {
        let sr = EMBED_SAMPLE_RATE;
        // Enroll the "matt" voice = 220 Hz tone.
        let embedder = FilterbankEmbedder::new();
        let enrolled_audio = tone_i16(220.0, 1.0, sr);
        let enrolled = Enrollment {
            label: "matt".into(),
            embedding: SpeakerEmbedding::new(embedder.embed(&enrolled_audio), embedder.model_id()),
        };
        let pipeline = SpeakerPipeline::baseline().with_enrollment(enrolled);

        // Mixed segment: enrolled tone + silence + different tone.
        let mut audio = tone_f32(220.0, 1.0, sr);
        audio.extend(silence_f32(0.5, sr));
        audio.extend(tone_f32(1_500.0, 1.0, sr));

        let turns = pipeline.classify_turns(&audio, sr).unwrap();
        assert_eq!(turns.len(), 2, "expected 2 turns, got {}", turns.len());
        assert_eq!(turns[0].role, TurnRole::EnrolledUser);
        assert!(matches!(turns[1].role, TurnRole::OtherSpeaker(_)));
        // Audio slices should be non-empty and roughly the right
        // shape (≈ 1s each at 16 kHz).
        assert!(turns[0].audio.len() > sr as usize / 2);
        assert!(turns[1].audio.len() > sr as usize / 2);
    }

    #[test]
    fn filter_to_enrolled_user_passthrough_without_enrollment() {
        let pipeline = SpeakerPipeline::baseline();
        let audio = tone_f32(220.0, 1.0, EMBED_SAMPLE_RATE);
        let filtered = pipeline
            .filter_to_enrolled_user(&audio, EMBED_SAMPLE_RATE)
            .unwrap()
            .unwrap();
        assert_eq!(filtered.len(), audio.len());
    }

    #[test]
    fn filter_to_enrolled_user_drops_non_matching() {
        let sr = EMBED_SAMPLE_RATE;
        let embedder = FilterbankEmbedder::new();
        let enrolled_audio = tone_i16(220.0, 1.0, sr);
        let enrolled = Enrollment {
            label: "matt".into(),
            embedding: SpeakerEmbedding::new(embedder.embed(&enrolled_audio), embedder.model_id()),
        };
        let pipeline = SpeakerPipeline::baseline().with_enrollment(enrolled);
        // Only non-matching audio.
        let other_only = tone_f32(1_800.0, 1.0, sr);
        let filtered = pipeline.filter_to_enrolled_user(&other_only, sr).unwrap();
        assert!(
            filtered.is_none(),
            "non-matching-only segment should return None"
        );
    }
}