mod sense_voice;

use crate::asr::sense_voice::SenseVoiceConfig;
use crate::error::*;
use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;
use tokio::sync::broadcast::{Receiver, Sender};
use tracing::debug;

/// 自动语音识别配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutomaticSpeechRecognitionConfig {
    pub text_channel_size: usize,
    pub sense_voice_config: SenseVoiceConfig,
}

/// 自动语音识别
pub struct AutomaticSpeechRecognition {
    sense_voice: sense_voice::SenseVoice,
    speech_receiver: Receiver<Vec<f32>>,
    text_sender: Sender<String>,
}

impl AutomaticSpeechRecognition {
    /// 初始化自动语音识别，传入语音接收实例以及识别结果发送实例
    pub fn init(
        config: AutomaticSpeechRecognitionConfig,
        speech_receiver: Receiver<Vec<f32>>,
    ) -> Result<Self> {
        debug!(
            "Initializing automatic speech recognition with config {:?}",
            &config
        );
        let sense_voice = sense_voice::SenseVoice::init(config.sense_voice_config.clone())?;
        let (text_sender, _) = broadcast::channel(config.text_channel_size);
        Ok(Self {
            sense_voice,
            speech_receiver,
            text_sender,
        })
    }
    
    pub fn get_text_receiver(&self) -> Receiver<String> {
        self.text_sender.subscribe()
    }

    /// 运行自动语音识别线程
    pub async fn run(&mut self) -> Result<()> {
        while let Ok(speech) = self.speech_receiver.recv().await {
            debug!("Received speech, len: {}", speech.len());
            let asr_result = self.sense_voice.infer(speech)?;
            debug!("Asr result: {:?}", asr_result);
            self.text_sender.send(asr_result)?;
        }
        Ok(())
    }
}

impl Default for AutomaticSpeechRecognitionConfig {
    fn default() -> Self {
        Self {
            text_channel_size: 5,
            sense_voice_config: SenseVoiceConfig::default(),
        }
    }
}