use std::{
    path::{Path, PathBuf},
    time::Duration,
};

use anyhow::{anyhow, Context, Result};
use ndarray::{Array1, Array2, Array3};
use ort::{session::Session, value::TensorRef};

#[derive(Debug, Clone)]
pub struct VadConfig {
    pub sample_rate: usize,
    pub threshold: f32,
    pub min_speech_ms: usize,
    pub min_silence_ms: usize,
    pub pad_ms: usize,
    pub timeout_seconds: u64,
}

impl Default for VadConfig {
    fn default() -> Self {
        Self {
            sample_rate: 16_000,
            threshold: 0.01,
            min_speech_ms: 180,
            min_silence_ms: 320,
            pad_ms: 220,
            timeout_seconds: 60,
        }
    }
}

pub struct VadEngine {
    model_path: Option<PathBuf>,
    config: VadConfig,
}

impl VadEngine {
    pub fn new(model_path: Option<String>, config: VadConfig) -> Result<Self> {
        let model_path = model_path.map(PathBuf::from);
        if let Some(path) = &model_path {
            if !path.exists() {
                return Err(anyhow!("vad model not found: {}", path.display()));
            }
        }

        Ok(Self { model_path, config })
    }

    pub async fn detect_segments(&self, samples: &[f32]) -> Vec<(f32, f32)> {
        if self.model_path.is_some() {
            let samples_owned = samples.to_vec();
            let model_path = self.model_path.clone().unwrap();
            let config = self.config.clone();
            let timeout_secs = self.config.timeout_seconds;

            match tokio::time::timeout(
                Duration::from_secs(timeout_secs),
                tokio::task::spawn_blocking(move || {
                    let mut session = match Self::load_onnx_session(&model_path) {
                        Ok(s) => s,
                        Err(e) => {
                            eprintln!("vad: failed to load onnx session: {e:#}");
                            return None;
                        }
                    };
                    Self::detect_with_onnx(&mut session, &samples_owned, &config).ok()
                }),
            )
            .await
            {
                Ok(Ok(Some(ranges))) if !ranges.is_empty() => {
                    eprintln!("vad: onnx detected {} speech ranges", ranges.len());
                    return ranges;
                }
                Ok(Ok(_)) => {}
                Ok(Err(e)) => {
                    eprintln!("vad: onnx error: {e:#}, falling back to energy detection");
                }
                Err(_) => {
                    eprintln!(
                        "vad: onnx timed out after {}s, falling back to energy detection",
                        timeout_secs
                    );
                }
            }
        }

        let ranges = self.detect_segments_with_energy(samples);
        eprintln!("vad: energy detection found {} speech ranges", ranges.len());
        ranges
    }

    fn load_onnx_session(model_path: &Path) -> Result<Session> {
        Session::builder()
            .context("failed to build onnx session")?
            .commit_from_file(model_path)
            .with_context(|| format!("failed to load silero vad model: {}", model_path.display()))
    }

    fn detect_with_onnx(
        session: &mut Session,
        samples: &[f32],
        config: &VadConfig,
    ) -> Result<Vec<(f32, f32)>> {
        let chunk_size = 512usize;
        let mut state = Array3::<f32>::zeros((2, 1, 128));
        let sr = Array1::<i64>::from_vec(vec![config.sample_rate as i64]);
        let mut speech_probabilities = Vec::new();

        for chunk in samples.chunks(chunk_size) {
            let mut padded = vec![0.0_f32; chunk_size];
            padded[..chunk.len()].copy_from_slice(chunk);
            let input = Array2::from_shape_vec((1, chunk_size), padded)
                .context("failed to build vad input tensor")?;

            let outputs = session
                .run(ort::inputs![
                    TensorRef::from_array_view(input.view())?,
                    TensorRef::from_array_view(sr.view())?,
                    TensorRef::from_array_view(state.view())?
                ])
                .context("silero vad inference failed")?;

            let first = &outputs[0];
            let (_, probs) = first
                .try_extract_tensor::<f32>()
                .context("failed to extract vad probabilities")?;
            let probability = probs.iter().copied().fold(0.0_f32, f32::max);
            speech_probabilities.push(probability);

            if outputs.len() > 1 {
                let second = &outputs[1];
                let (_, next_state) = second
                    .try_extract_tensor::<f32>()
                    .context("failed to extract vad state")?;
                if next_state.len() == state.len() {
                    state = Array3::from_shape_vec((2, 1, 128), next_state.to_vec())
                        .context("failed to rebuild vad state")?;
                }
            }
        }

        Ok(Self::merge_probabilities(&speech_probabilities, chunk_size, config))
    }

    fn detect_segments_with_energy(&self, samples: &[f32]) -> Vec<(f32, f32)> {
        let frame_size = (self.config.sample_rate / 50).max(1);
        let mut energies = Vec::new();
        for chunk in samples.chunks(frame_size) {
            let energy = chunk.iter().map(|sample| sample.abs()).sum::<f32>() / chunk.len() as f32;
            energies.push(energy);
        }

        if energies.is_empty() {
            return Vec::new();
        }

        let dynamic_threshold = self.dynamic_energy_threshold(&energies);
        eprintln!(
            "vad: using energy fallback, frames={}, threshold={:.5}",
            energies.len(),
            dynamic_threshold
        );
        Self::merge_probabilities_with_threshold(&energies, frame_size, dynamic_threshold, &self.config)
    }

    fn merge_probabilities(frames: &[f32], frame_size: usize, config: &VadConfig) -> Vec<(f32, f32)> {
        Self::merge_probabilities_with_threshold(frames, frame_size, config.threshold, config)
    }

    fn merge_probabilities_with_threshold(
        frames: &[f32],
        frame_size: usize,
        threshold: f32,
        config: &VadConfig,
    ) -> Vec<(f32, f32)> {
        let min_speech_frames = (config.min_speech_ms / 20).max(1);
        let min_silence_frames = (config.min_silence_ms / 20).max(1);
        let pad_seconds = config.pad_ms as f32 / 1000.0;

        let mut result = Vec::new();
        let mut start_frame: Option<usize> = None;
        let mut silent_frames = 0usize;

        for (index, probability) in frames.iter().enumerate() {
            if *probability >= threshold {
                if start_frame.is_none() {
                    start_frame = Some(index);
                }
                silent_frames = 0;
                continue;
            }

            if let Some(start) = start_frame {
                silent_frames += 1;
                if silent_frames >= min_silence_frames {
                    let end_frame = index.saturating_sub(silent_frames);
                    if end_frame.saturating_sub(start) >= min_speech_frames {
                        let start_sec = (start * frame_size) as f32 / config.sample_rate as f32;
                        let end_sec = ((end_frame + 1) * frame_size) as f32 / config.sample_rate as f32;
                        result.push(((start_sec - pad_seconds).max(0.0), end_sec + pad_seconds));
                    }
                    start_frame = None;
                    silent_frames = 0;
                }
            }
        }

        if let Some(start) = start_frame {
            let end_frame = frames.len().saturating_sub(1);
            if end_frame.saturating_sub(start) >= min_speech_frames {
                let start_sec = (start * frame_size) as f32 / config.sample_rate as f32;
                let end_sec = ((end_frame + 1) * frame_size) as f32 / config.sample_rate as f32;
                result.push(((start_sec - pad_seconds).max(0.0), end_sec + pad_seconds));
            }
        }

        if result.is_empty() && !frames.is_empty() {
            let total_seconds = (frames.len() * frame_size) as f32 / config.sample_rate as f32;
            result.push((0.0, total_seconds));
        }

        result
    }

    fn dynamic_energy_threshold(&self, energies: &[f32]) -> f32 {
        let mut sorted = energies.to_vec();
        sorted.sort_by(|left, right| left.total_cmp(right));

        let mean = energies.iter().sum::<f32>() / energies.len() as f32;
        let p70 = percentile(&sorted, 0.70);
        let p90 = percentile(&sorted, 0.90);
        let adaptive = (mean + (p90 - mean) * 0.18).max(p70 * 0.72);

        adaptive.max(self.config.threshold * 0.45)
    }
}

fn percentile(sorted: &[f32], ratio: f32) -> f32 {
    if sorted.is_empty() {
        return 0.0;
    }

    let index = ((sorted.len() - 1) as f32 * ratio.clamp(0.0, 1.0)).round() as usize;
    sorted[index.min(sorted.len() - 1)]
}
