J.A.R.V.I.S-rust/crates/jarvis-core/src/stt/vosk.rs

use once_cell::sync::OnceCell;
use vosk::{DecodingState, Model, Recognizer};

use std::sync::Mutex;

// use crate::config::VOSK_MODEL_PATH;
use crate::config;
use crate::stt::vosk_models;
use crate::DB;

static MODEL: OnceCell<Model> = OnceCell::new();
static RECOGNIZER: OnceCell<Mutex<Recognizer>> = OnceCell::new();

pub fn init_vosk() -> Result<(), String> {
    if RECOGNIZER.get().is_some() {
        return Ok(());
    } // already initialized

    let model_path = get_configured_model_path()?;

    let model = Model::new(model_path.to_str().unwrap())
        .ok_or_else(|| format!("Failed to load Vosk model from: {}", model_path.display()))?;

    let mut recognizer = Recognizer::new(&model, 16000.0)
        .ok_or("Failed to create Vosk recognizer")?;

    recognizer.set_max_alternatives(10);
    recognizer.set_words(true);
    recognizer.set_partial_words(true);

    MODEL.set(model);
    RECOGNIZER.set(Mutex::new(recognizer));

    Ok(())
}

pub fn recognize(data: &[i16], include_partial: bool) -> Option<String> {
    let state = RECOGNIZER
        .get()
        .unwrap()
        .lock()
        .unwrap()
        .accept_waveform(data);

    match state {
        Ok(ds) => {
            match ds {
                DecodingState::Running => {
                    if include_partial {
                        Some(
                            RECOGNIZER
                                .get()
                                .unwrap()
                                .lock()
                                .unwrap()
                                .partial_result()
                                .partial
                                .into(),
                        )
                    } else {
                        None
                    }
                }
                DecodingState::Finalized => {
                    // Result will always be multiple because we called set_max_alternatives
                    RECOGNIZER
                        .get()
                        .unwrap()
                        .lock()
                        .unwrap()
                        .result()
                        .multiple()
                        .and_then(|m| m.alternatives.first().map(|a| a.text.to_string()))
                }
                DecodingState::Failed => None,
            }
        },
        Err(err) => {
                error!("Vosk accept waveform error.\nError details: {}", err);

                None
        }
    }
}

fn get_configured_model_path() -> Result<std::path::PathBuf, String> {
    // try to get from settings
    if let Some(db) = DB.get() {
        let settings = db.read();
        if !settings.vosk_model.is_empty() {
            if let Some(path) = vosk_models::get_model_path(&settings.vosk_model) {
                return Ok(path);
            }
            warn!("Configured Vosk model '{}' not found, falling back to auto-detect", settings.vosk_model);
        }
    }
    
    // auto-detect: use first available model
    let available = vosk_models::scan_vosk_models();
    if let Some(first) = available.first() {
        info!("Auto-detected Vosk model: {}", first.name);
        return Ok(first.path.clone());
    }
    
    // fallback to legacy path
    let legacy_path = std::path::Path::new(config::VOSK_MODEL_PATH);
    if legacy_path.exists() {
        return Ok(legacy_path.to_path_buf());
    }
    
    Err("No Vosk models found".into())
}

// pub fn stereo_to_mono(input_data: &[i16]) -> Vec<i16> {
//     let mut result = Vec::with_capacity(input_data.len() / 2);
//     result.extend(
//         input_data
//             .chunks_exact(2)
//             .map(|chunk| chunk[0] / 2 + chunk[1] / 2),
//     );

//     result
// }