//! Silero VAD v6 ONNX runtime integration (P1 VAD_002).
//!
//! ## Silero VAD v6 model I/O
//!
//! The v6 model (silero_vad.onnx from the v6.x releases) has a different
//! interface from v4. Key changes:
//!
//! | Tensor  | Shape            | Dtype | Meaning                              |
//! |---------|------------------|-------|--------------------------------------|
//! | input   | \[1, 576\]       | f32   | 64-sample context + 512-sample frame |
//! | state   | \[2, 1, 128\]    | f32   | LSTM state (carry across frames)     |
//! | sr      | \[1\]            | i64   | Sample rate (16000 or 8000)          |
//! | output  | \[1, 1\]         | f32   | Speech probability                   |
//! | stateN  | \[2, 1, 128\]    | f32   | Updated LSTM state                   |
//!
//! Frame size: **512 samples at 16 kHz = 32 ms**.
//! Context: **64 samples** prepended to each frame (last 64 samples of previous frame).
//! Total input width: 512 + 64 = **576 samples**.
//!
//! ## Threading
//!
//! `SileroOnnxVad` is `Send`. The session is created once and reused —
//! never re-created per callback (INV_007).
//!
//! ## Accumulation
//!
//! The capture pipeline delivers 10 ms frames (480 samples at 48 kHz →
//! 160 samples at 16 kHz). Three 10 ms frames = 30 ms ≈ 32 ms. We
//! accumulate 512 samples (32 ms at 16 kHz) before running inference.
//! The last probability is held between inference calls so the state
//! machine always has a value to work with.
//!
//! ## Fallback
//!
//! `try_new` returns `None` when the model file is missing, the ONNX
//! Runtime is unavailable, or the platform is not iOS/macOS. The caller
//! falls back to `WebRtcFallbackVad`.

use super::{VadOutput, VoiceActivityDetector};
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::thread::JoinHandle;

/// 16 kHz frame size for Silero VAD v6 (32 ms).
pub const SILERO_FRAME_16K: usize = 512;
/// Context size prepended to each frame (64 samples at 16 kHz).
pub const SILERO_CONTEXT_16K: usize = 64;
/// Total input width: context + frame.
pub const SILERO_INPUT_WIDTH: usize = SILERO_CONTEXT_16K + SILERO_FRAME_16K;
/// LSTM state size: 2 × 1 × 128 = 256 f32 values.
pub const SILERO_STATE_SIZE: usize = 256;
/// Maximum lag in 10 ms frames before the realtime callback treats
/// the Silero worker as stale and falls back to the local WebRTC
/// detector for that frame.
pub const SILERO_MAX_STALE_FRAMES: u64 = 3;

/// Silero VAD v6 ONNX backend.
///
/// Operates at **16 kHz**, accumulating 32 ms frames (512 samples)
/// before running inference. The caller is responsible for downsampling
/// from 48 kHz before calling `process_10ms`.
pub struct SileroOnnxVad {
    /// LSTM state [2, 1, 128] — persisted across frames.
    state: Box<[f32; SILERO_STATE_SIZE]>,
    /// Context ring: last 64 samples of the previous frame.
    context: Box<[f32; SILERO_CONTEXT_16K]>,
    /// Accumulation buffer for 16 kHz samples (fills to SILERO_FRAME_16K).
    accum: Vec<f32>,
    /// Last speech probability output (held between inference calls).
    last_probability: f32,
    /// Model path stored for diagnostics.
    model_path: String,
    /// Inner ONNX implementation (platform-specific).
    inner: SileroInner,
}

enum SileroInner {
    Onnx(OnnxSession),
    #[cfg(test)]
    Stub,
}

struct OnnxSession {
    session: ort::session::Session,
}

impl SileroOnnxVad {
    /// Attempt to load the Silero v6 ONNX model from `model_path`.
    ///
    /// Returns `None` when the model file is missing, the ONNX Runtime
    /// is unavailable, or the platform does not support ONNX.
    pub fn try_new(model_path: &str) -> Option<Self> {
        Self::try_new_onnx(model_path)
    }

    fn try_new_onnx(model_path: &str) -> Option<Self> {
        use tracing::{error, info};

        if !std::path::Path::new(model_path).exists() {
            tracing::warn!(
                target: "chanora_audio",
                path = model_path,
                "SileroOnnxVad: model file not found; falling back to WebRtcFallbackVad"
            );
            return None;
        }

        let session_result = std::panic::catch_unwind(|| {
            ort::session::Session::builder().and_then(|mut b| b.commit_from_file(model_path))
        });

        match session_result {
            Err(_) => {
                error!(
                    target: "chanora_audio",
                    path = model_path,
                    "SileroOnnxVad: ONNX Runtime panicked during load; falling back to WebRtcFallbackVad"
                );
                None
            }
            Ok(Ok(session)) => {
                info!(
                    target: "chanora_audio",
                    path = model_path,
                    "SileroOnnxVad v6: model loaded"
                );
                Some(Self {
                    state: Box::new([0.0; SILERO_STATE_SIZE]),
                    context: Box::new([0.0; SILERO_CONTEXT_16K]),
                    accum: Vec::with_capacity(SILERO_FRAME_16K),
                    last_probability: 0.0,
                    model_path: model_path.to_owned(),
                    inner: SileroInner::Onnx(OnnxSession { session }),
                })
            }
            Ok(Err(e)) => {
                error!(
                    target: "chanora_audio",
                    path = model_path,
                    error = %e,
                    "SileroOnnxVad: failed to load model; falling back to WebRtcFallbackVad"
                );
                None
            }
        }
    }

    /// Reset LSTM state and context (call on voice_leave / session restart).
    pub fn reset_state(&mut self) {
        self.state.iter_mut().for_each(|v| *v = 0.0);
        self.context.iter_mut().for_each(|v| *v = 0.0);
        self.accum.clear();
        self.last_probability = 0.0;
    }

    /// Return the model path for diagnostics.
    pub fn model_path(&self) -> &str {
        &self.model_path
    }

    fn input_with_context(context: &[f32; SILERO_CONTEXT_16K], audio_frame: &[f32]) -> Vec<f32> {
        let mut input = Vec::with_capacity(SILERO_CONTEXT_16K + audio_frame.len());
        input.extend_from_slice(context);
        input.extend_from_slice(audio_frame);
        input
    }

    fn update_context_from_frame(&mut self, audio_frame: &[f32]) {
        let ctx_start = audio_frame.len().saturating_sub(SILERO_CONTEXT_16K);
        let new_ctx = &audio_frame[ctx_start..];
        let copy_len = new_ctx.len().min(SILERO_CONTEXT_16K);
        self.context.fill(0.0);
        self.context[SILERO_CONTEXT_16K - copy_len..].copy_from_slice(&new_ctx[..copy_len]);
    }

    /// Run one upstream-style `calc_level` pass over a 32 ms / 512-sample
    /// 16 kHz frame: concatenate prior context, pass `input/state/sr` to
    /// ONNX, persist `stateN`, then refresh context from the current frame.
    /// Updates `last_probability` and returns the new value.
    fn calc_level(&mut self, audio_frame: &[f32]) -> f32 {
        use ort::value::Value;
        use tracing::error;

        #[cfg(test)]
        let inner = match self.inner {
            SileroInner::Onnx(ref mut inner) => inner,
            SileroInner::Stub => return self.last_probability,
        };
        #[cfg(not(test))]
        let SileroInner::Onnx(ref mut inner) = self.inner;

        debug_assert_eq!(audio_frame.len(), SILERO_FRAME_16K);

        // Build input: [1, 576] = context (64) + frame (512), matching
        // snakers4/silero-vad's Rust `calc_level` example.
        let input_vec = Self::input_with_context(self.context.as_ref(), audio_frame);

        // Build ndarray tensors.
        use ndarray::{Array, IxDyn};

        let input_arr = Array::from_shape_vec(IxDyn(&[1, SILERO_INPUT_WIDTH]), input_vec);
        let state_arr = Array::from_shape_vec(IxDyn(&[2, 1, 128]), self.state.to_vec());
        let sr_arr = Array::from_shape_vec(IxDyn(&[1]), vec![16000_i64]);

        let (input_arr, state_arr, sr_arr) = match (input_arr, state_arr, sr_arr) {
            (Ok(i), Ok(s), Ok(sr)) => (i, s, sr),
            _ => return self.last_probability,
        };

        let input_val = match Value::from_array(input_arr) {
            Ok(v) => v,
            Err(e) => {
                error!(target: "chanora_audio", error = %e, "SileroOnnxVad: input tensor error");
                return self.last_probability;
            }
        };
        let state_val = match Value::from_array(state_arr) {
            Ok(v) => v,
            Err(e) => {
                error!(target: "chanora_audio", error = %e, "SileroOnnxVad: state tensor error");
                return self.last_probability;
            }
        };
        let sr_val = match Value::from_array(sr_arr) {
            Ok(v) => v,
            Err(e) => {
                error!(target: "chanora_audio", error = %e, "SileroOnnxVad: sr tensor error");
                return self.last_probability;
            }
        };

        let outputs =
            match inner
                .session
                .run([(&input_val).into(), (&state_val).into(), (&sr_val).into()])
            {
                Ok(o) => o,
                Err(e) => {
                    error!(target: "chanora_audio", error = %e, "SileroOnnxVad: inference failed");
                    return self.last_probability;
                }
            };

        // Extract probability from "output".
        if let Ok((_, prob_data)) = outputs["output"].try_extract_tensor::<f32>() {
            if let Some(&p) = prob_data.first() {
                self.last_probability = p.clamp(0.0, 1.0);
            }
        }

        // Update state from "stateN".
        if let Ok((shape, state_data)) = outputs["stateN"].try_extract_tensor::<f32>() {
            let total: usize = shape.iter().map(|&d| d as usize).product();
            let copy_len = total.min(SILERO_STATE_SIZE);
            self.state[..copy_len].copy_from_slice(&state_data[..copy_len]);
        }

        drop(outputs);

        // Match the upstream example: context becomes the last context_size
        // samples from the current frame after the model call succeeds.
        self.update_context_from_frame(audio_frame);

        self.last_probability
    }
}

impl VoiceActivityDetector for SileroOnnxVad {
    /// Accept one 10 ms **16 kHz** f32 mono frame (160 samples).
    ///
    /// Accumulates samples until a full 32 ms frame (512 samples) is
    /// ready, then runs inference. Between inference calls the last
    /// probability is returned unchanged.
    fn process_10ms(&mut self, samples: &[f32]) -> VadOutput {
        debug_assert_eq!(
            samples.len(),
            super::resampler::OUTPUT_FRAME_10MS,
            "SileroOnnxVad expects 160 samples (16 kHz 10 ms), got {}",
            samples.len()
        );

        self.accum.extend_from_slice(samples);

        if self.accum.len() >= SILERO_FRAME_16K {
            let audio_frame: Vec<f32> = self.accum[..SILERO_FRAME_16K].to_vec();
            if matches!(self.inner, SileroInner::Onnx(_)) {
                self.calc_level(&audio_frame);
            } else {
                self.update_context_from_frame(&audio_frame);
            }
            // Drain the accumulator (keep any overflow for next frame).
            let overflow: Vec<f32> = self.accum.drain(SILERO_FRAME_16K..).collect();
            self.accum.clear();
            self.accum.extend_from_slice(&overflow);
        }

        VadOutput {
            probability: self.last_probability,
            speech: self.last_probability >= 0.5,
        }
    }
}

// SAFETY: ONNX Runtime sessions are thread-safe for inference.
// State arrays are owned by this struct and accessed only from
// the single capture callback thread.
unsafe impl Send for SileroOnnxVad {}

struct SileroFrameMessage {
    seq: u64,
    frame: [f32; super::resampler::INPUT_FRAME_10MS],
}

/// Background Silero worker. The realtime callback only enqueues
/// 10 ms frames and reads the latest probability atomically.
pub struct SileroOnnxVadWorker {
    tx: Option<std::sync::mpsc::SyncSender<SileroFrameMessage>>,
    latest_probability: Arc<AtomicU32>,
    latest_processed_seq: Arc<AtomicU64>,
    alive: Arc<AtomicBool>,
    handle: Option<JoinHandle<()>>,
}

impl SileroOnnxVadWorker {
    /// Start a background Silero worker if the model loads.
    pub fn try_new(model_path: &str) -> Option<Self> {
        let vad = SileroOnnxVad::try_new(model_path)?;
        let latest_probability = Arc::new(AtomicU32::new(0.0_f32.to_bits()));
        let latest_processed_seq = Arc::new(AtomicU64::new(u64::MAX));
        let alive = Arc::new(AtomicBool::new(true));
        let (tx, rx) = std::sync::mpsc::sync_channel::<SileroFrameMessage>(64);
        let latest_probability_for_thread = latest_probability.clone();
        let latest_processed_seq_for_thread = latest_processed_seq.clone();
        let alive_for_thread = alive.clone();

        let handle = std::thread::Builder::new()
            .name("chanora-silero-vad".to_string())
            .spawn(move || {
                let mut vad = super::Resampled16kHzVad::new(vad);
                while alive_for_thread.load(Ordering::Relaxed) {
                    let message = match rx.recv() {
                        Ok(message) => message,
                        Err(_) => break,
                    };
                    let output = vad.process_10ms(&message.frame);
                    latest_probability_for_thread.store(
                        output.probability.clamp(0.0, 1.0).to_bits(),
                        Ordering::Relaxed,
                    );
                    latest_processed_seq_for_thread.store(message.seq, Ordering::Relaxed);
                }
            })
            .ok()?;

        Some(Self {
            tx: Some(tx),
            latest_probability,
            latest_processed_seq,
            alive,
            handle: Some(handle),
        })
    }

    /// Best-effort enqueue of a 10 ms frame for background inference.
    pub fn try_send(&self, seq: u64, frame: &[f32; super::resampler::INPUT_FRAME_10MS]) -> bool {
        let Some(tx) = &self.tx else {
            return false;
        };
        tx.try_send(SileroFrameMessage { seq, frame: *frame })
            .is_ok()
    }

    /// Latest probability published by the background worker.
    pub fn latest_probability(&self) -> f32 {
        f32::from_bits(self.latest_probability.load(Ordering::Relaxed))
    }

    /// Number of 10 ms frames the worker is behind the capture thread.
    pub fn lag_frames(&self, capture_seq: u64) -> u64 {
        capture_seq.saturating_sub(self.latest_processed_seq.load(Ordering::Relaxed))
    }

    /// True when the worker is too far behind to trust its latest
    /// probability for the current frame.
    pub fn is_stale(&self, capture_seq: u64) -> bool {
        let latest = self
            .latest_processed_seq
            .load(std::sync::atomic::Ordering::Relaxed);
        latest == u64::MAX || capture_seq.saturating_sub(latest) > SILERO_MAX_STALE_FRAMES
    }
}

impl Drop for SileroOnnxVadWorker {
    fn drop(&mut self) {
        self.alive.store(false, Ordering::Relaxed);
        let _ = self.tx.take();
        let _ = self.handle.take();
    }
}

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

    fn make_stub_vad() -> SileroOnnxVad {
        SileroOnnxVad {
            state: Box::new([0.0; SILERO_STATE_SIZE]),
            context: Box::new([0.0; SILERO_CONTEXT_16K]),
            accum: Vec::new(),
            last_probability: 0.0,
            model_path: String::new(),
            inner: SileroInner::Stub,
        }
    }

    #[test]
    fn try_new_returns_none_without_model_file() {
        let result = SileroOnnxVad::try_new("/nonexistent/silero_vad.onnx");
        assert!(result.is_none());
    }

    #[test]
    fn stub_accumulates_and_holds_zero_probability() {
        let mut vad = make_stub_vad();
        let frame = vec![0.0_f32; super::super::resampler::OUTPUT_FRAME_10MS];
        // Feed 3 frames (30 ms < 32 ms) — no inference yet.
        for _ in 0..3 {
            let out = VoiceActivityDetector::process_10ms(&mut vad, &frame);
            assert_eq!(out.probability, 0.0);
        }
        // Feed 1 more frame (40 ms > 32 ms) — accumulator drains.
        let out = VoiceActivityDetector::process_10ms(&mut vad, &frame);
        assert_eq!(out.probability, 0.0); // stub stays at 0
    }

    #[test]
    fn reset_state_clears_all() {
        let mut vad = make_stub_vad();
        vad.state[0] = 1.0;
        vad.context[0] = 1.0;
        vad.last_probability = 0.9;
        vad.accum.push(0.5);
        vad.reset_state();
        assert_eq!(vad.state[0], 0.0);
        assert_eq!(vad.context[0], 0.0);
        assert_eq!(vad.last_probability, 0.0);
        assert!(vad.accum.is_empty());
    }

    #[test]
    fn accumulates_correct_number_of_samples() {
        let mut vad = make_stub_vad();
        let frame = vec![0.1_f32; super::super::resampler::OUTPUT_FRAME_10MS]; // 160 samples
                                                                               // 3 × 160 = 480 < 512 — not yet full.
        for _ in 0..3 {
            VoiceActivityDetector::process_10ms(&mut vad, &frame);
        }
        assert_eq!(vad.accum.len(), 480);
        // 4th frame: 640 > 512 — inference fires, 128 samples remain.
        VoiceActivityDetector::process_10ms(&mut vad, &frame);
        assert_eq!(vad.accum.len(), 128);
    }

    #[test]
    fn input_concatenates_context_before_frame_like_upstream_example() {
        let mut context = [0.0_f32; SILERO_CONTEXT_16K];
        context[0] = -1.0;
        context[SILERO_CONTEXT_16K - 1] = 1.0;
        let frame = vec![0.25_f32; SILERO_FRAME_16K];

        let input = SileroOnnxVad::input_with_context(&context, &frame);

        assert_eq!(input.len(), SILERO_INPUT_WIDTH);
        assert_eq!(input[0], -1.0);
        assert_eq!(input[SILERO_CONTEXT_16K - 1], 1.0);
        assert_eq!(input[SILERO_CONTEXT_16K], 0.25);
        assert_eq!(input[SILERO_INPUT_WIDTH - 1], 0.25);
    }

    #[test]
    fn context_tracks_last_64_samples_of_completed_frame() {
        let mut vad = make_stub_vad();
        let frame = vec![0.0_f32; super::super::resampler::OUTPUT_FRAME_10MS];
        for idx in 0..4 {
            let mut chunk = frame.clone();
            let chunk_len = chunk.len();
            for (sample_idx, sample) in chunk.iter_mut().enumerate() {
                *sample = (idx * chunk_len + sample_idx) as f32;
            }
            VoiceActivityDetector::process_10ms(&mut vad, &chunk);
        }

        let completed_frame: Vec<f32> = (0..SILERO_FRAME_16K).map(|v| v as f32).collect();
        assert_eq!(
            vad.context.as_ref(),
            &completed_frame[SILERO_FRAME_16K - SILERO_CONTEXT_16K..]
        );
    }
}