Lumotia/crates/transcription/src/local_engine.rs

use std::path::Path;
use std::sync::Mutex;
use std::time::Instant;

use transcribe_rs::{SpeechModel, TranscribeOptions, TranscriptionResult};

use kon_core::error::{KonError, Result};
use kon_core::types::{
    AudioSamples, EngineName, ModelId, Segment, Transcript, TranscriptionOptions,
};

use crate::whisper_rs_backend::WhisperRsBackend;

/// Result of a timed transcription: transcript + inference duration.
pub struct TimedTranscript {
    pub transcript: Transcript,
    pub inference_ms: u64,
}

/// Public discriminator selected by the loaders (`load_parakeet`, `load_whisper`)
/// and passed to `LocalEngine::load`. `src-tauri::commands::models` names this
/// type as the return of `load_model_from_disk`, so it must be `pub`.
pub enum SpeechBackend {
    /// transcribe-rs-owned model. Used for Parakeet ONNX (wrapped in
    /// ParakeetWordGranularity for word-level timestamps).
    Adapter(Box<dyn SpeechModel + Send>),
    /// Direct whisper-rs. The only path that actually forwards `initial_prompt`.
    WhisperRs(WhisperRsBackend),
}

/// Wraps any transcribe-rs engine in Kon's SpeechToText trait.
/// Encapsulates threading: inference always runs on a blocking thread.
/// The rest of the app never imports transcribe-rs directly.
pub struct LocalEngine {
    engine: Mutex<Option<SpeechBackend>>,
    engine_name: EngineName,
    loaded_model_id: Mutex<Option<ModelId>>,
}

impl LocalEngine {
    pub fn new(engine_name: EngineName) -> Self {
        Self {
            engine: Mutex::new(None),
            engine_name,
            loaded_model_id: Mutex::new(None),
        }
    }

    pub fn load(&self, backend: SpeechBackend, model_id: ModelId) {
        let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
        *guard = Some(backend);
        let mut id_guard = self
            .loaded_model_id
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        *id_guard = Some(model_id);
    }

    pub fn name(&self) -> &EngineName {
        &self.engine_name
    }

    pub fn loaded_model_id(&self) -> Option<ModelId> {
        let guard = self
            .loaded_model_id
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        guard.clone()
    }

    pub fn is_loaded(&self) -> bool {
        let guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
        guard.is_some()
    }

    /// Run transcription synchronously with timing.
    /// Called from within spawn_blocking.
    pub fn transcribe_sync(
        &self,
        audio: &AudioSamples,
        options: &TranscriptionOptions,
    ) -> Result<TimedTranscript> {
        let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
        let backend = guard.as_mut().ok_or(KonError::EngineNotLoaded)?;

        let start = Instant::now();
        let segments: Vec<Segment> = match backend {
            SpeechBackend::Adapter(model) => {
                let opts = TranscribeOptions {
                    language: options.language.clone(),
                    translate: false,
                    leading_silence_ms: None,
                    trailing_silence_ms: None,
                };
                let result: TranscriptionResult = model
                    .transcribe(audio.samples(), &opts)
                    .map_err(|e| KonError::TranscriptionFailed(e.to_string()))?;
                result
                    .segments
                    .unwrap_or_default()
                    .into_iter()
                    .map(|s| Segment {
                        start: s.start as f64,
                        end: s.end as f64,
                        text: s.text,
                    })
                    .collect()
            }
            SpeechBackend::WhisperRs(w) => w
                .transcribe_sync(audio.samples(), options)
                .map_err(|e| KonError::TranscriptionFailed(e.to_string()))?,
        };
        let inference_ms = start.elapsed().as_millis() as u64;

        Ok(TimedTranscript {
            transcript: Transcript::new(
                segments,
                options.language.clone().unwrap_or_else(|| "en".to_string()),
                audio.duration_secs(),
            ),
            inference_ms,
        })
    }
}

/// Thin wrapper over `ParakeetModel` that overrides `transcribe_raw` to
/// request word-granularity segments. `transcribe-rs` 0.3's trait impl for
/// `ParakeetModel::transcribe_raw` ignores `TranscribeOptions` and uses
/// `TimestampGranularity::Token` (per-subword) — which surfaces in Kon as
/// "T Est Ing . One , Two , Three" output. The concrete-type method
/// `ParakeetModel::transcribe_with` accepts `ParakeetParams` with an
/// explicit granularity; this wrapper exposes that to the trait object.
struct ParakeetWordGranularity(transcribe_rs::onnx::parakeet::ParakeetModel);

impl transcribe_rs::SpeechModel for ParakeetWordGranularity {
    fn capabilities(&self) -> transcribe_rs::ModelCapabilities {
        self.0.capabilities()
    }

    fn default_leading_silence_ms(&self) -> u32 {
        self.0.default_leading_silence_ms()
    }

    fn default_trailing_silence_ms(&self) -> u32 {
        self.0.default_trailing_silence_ms()
    }

    fn transcribe_raw(
        &mut self,
        samples: &[f32],
        options: &TranscribeOptions,
    ) -> std::result::Result<TranscriptionResult, transcribe_rs::TranscribeError> {
        use transcribe_rs::onnx::parakeet::{ParakeetParams, TimestampGranularity};
        let params = ParakeetParams {
            language: options.language.clone(),
            timestamp_granularity: Some(TimestampGranularity::Word),
        };
        self.0.transcribe_with(samples, &params)
    }
}

/// Load a Parakeet model from a directory path.
pub fn load_parakeet(model_dir: &Path) -> Result<SpeechBackend> {
    use transcribe_rs::onnx::Quantization;
    let model = transcribe_rs::onnx::parakeet::ParakeetModel::load(model_dir, &Quantization::Int8)
        .map_err(|e| KonError::TranscriptionFailed(format!("Failed to load Parakeet: {e}")))?;
    Ok(SpeechBackend::Adapter(Box::new(ParakeetWordGranularity(
        model,
    ))))
}

/// Load a Whisper model from a GGML file path via whisper-rs.
pub fn load_whisper(model_path: &Path) -> Result<SpeechBackend> {
    let backend = WhisperRsBackend::load(model_path)
        .map_err(|e| KonError::TranscriptionFailed(format!("Failed to load Whisper: {e}")))?;
    Ok(SpeechBackend::WhisperRs(backend))
}

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

    #[test]
    fn engine_reports_not_available_before_loading() {
        let engine = LocalEngine::new(EngineName::new("test"));
        assert!(!engine.is_loaded());
        assert!(engine.loaded_model_id().is_none());
    }
}