use crate::error::*;
use ndarray::{Array2, ArrayD};
use ort::session::Session;
use ort::session::builder::GraphOptimizationLevel;
use ort::value::Tensor;
use serde::{Deserialize, Serialize};

/// Silero语音活动检测模型配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SileroConfig {
    pub sample_rate: usize,
    pub chunk_size: usize,
    pub model_path: String,
}

/// Silero语音活动检测模型
pub struct Silero {
    pub(crate) config: SileroConfig,
    state: ArrayD<f32>,
    session: Session,
}

impl Silero {
    /// 使用配置初始化Silero语音活动检测模型
    pub fn init(mut config: SileroConfig) -> Result<Self> {
        let chunk_size = match &config.sample_rate {
            8000 => 256,
            16000 => 512,
            _ => {
                return Err(NihilityListenerError::Init(format!(
                    "Sample rate {} not supported",
                    config.sample_rate
                )));
            }
        };
        if config.chunk_size != chunk_size {
            config.chunk_size = chunk_size;
        }
        let session = Session::builder()?
            .with_optimization_level(GraphOptimizationLevel::Level3)?
            .with_inter_threads(1)?
            .with_intra_threads(1)?
            .commit_from_file(&config.model_path)?;
        Ok(Self {
            config,
            state: ndarray::Array3::<f32>::zeros((2, 1, 128)).into_dyn(),
            session,
        })
    }

    /// 预测当前语音数据块语音活动概率
    pub fn predict(&mut self, samples: &[f32]) -> Result<f32> {
        if samples.len() != self.config.chunk_size {
            return Err(NihilityListenerError::Unknown(format!(
                "samples len must be {}",
                self.config.chunk_size
            )));
        }
        let input = Array2::from_shape_vec((1, self.config.chunk_size), samples.to_vec())?;
        let sample_rate = ndarray::arr0::<i64>(self.config.sample_rate as i64);

        let state_taken = std::mem::take(&mut self.state);

        // ort::inputs! macro no longer supports ArrayView, use Tensor::from_array instead
        let inputs = ort::inputs![
            Tensor::from_array(input.to_owned())?,
            Tensor::from_array(state_taken.to_owned())?,
            Tensor::from_array(sample_rate.to_owned())?,
        ];

        let outputs = self.session.run(inputs)?;

        // 递归地更新状态。
        self.state = outputs
            .get("stateN")
            .expect("model error, output do not have stateN field")
            .try_extract_array::<f32>()?
            .to_owned();

        // 获取语音发生的概率。
        let output = outputs
            .get("output")
            .expect("model error, output do not have output field")
            .try_extract_array::<f32>()?;
        Ok(output[[0, 0]])
    }
}

impl Default for SileroConfig {
    fn default() -> Self {
        Self {
            sample_rate: 16000,
            chunk_size: 512,
            model_path: "model/silero_vad.onnx".to_string(),
        }
    }
}
