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::transcriber::{Transcriber, TranscriberCapabilities};
#[cfg(feature = "whisper")]
use crate::whisper_rs_backend::WhisperRsBackend;

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

/// Adapts any `transcribe-rs` `SpeechModel` into the `Transcriber`
/// trait. Today this is only used for Parakeet (ONNX), but the adapter
/// is the path any future transcribe-rs-backed engine plugs through —
/// Moonshine, fine-tuned Parakeet variants, etc.
pub struct SpeechModelAdapter(pub Box<dyn SpeechModel + Send>);

impl Transcriber for SpeechModelAdapter {
    fn capabilities(&self) -> TranscriberCapabilities {
        TranscriberCapabilities {
            sample_rate: kon_core::constants::WHISPER_SAMPLE_RATE,
            channels: 1,
            supports_initial_prompt: false,
        }
    }

    fn transcribe_sync(
        &mut self,
        samples: &[f32],
        options: &TranscriptionOptions,
    ) -> Result<Vec<Segment>> {
        let opts = TranscribeOptions {
            language: options.language.clone(),
            translate: false,
            leading_silence_ms: None,
            trailing_silence_ms: None,
        };
        let result: TranscriptionResult = self
            .0
            .transcribe(samples, &opts)
            .map_err(|e| KonError::TranscriptionFailed(e.to_string()))?;
        Ok(result
            .segments
            .unwrap_or_default()
            .into_iter()
            .map(|s| Segment {
                start: s.start as f64,
                end: s.end as f64,
                text: s.text,
            })
            .collect())
    }
}

/// Owns the currently-loaded speech backend and serialises inference
/// against model-swap operations via a `Mutex`. All transcription goes
/// through this struct; no caller ever holds a raw `Box<dyn Transcriber>`.
pub struct LocalEngine {
    engine: Mutex<Option<Box<dyn Transcriber + Send>>>,
    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: Box<dyn Transcriber + Send>, 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);
    }

    /// Drop the loaded model and free its backing resources (GPU VRAM,
    /// CPU memory, mmap'd GGML tensors). Used by the sequential-GPU
    /// guard (brief item A.1 #28) so loading the LLM on a tight-VRAM
    /// system first frees the transcription engine, and vice versa.
    ///
    /// No-op when nothing is loaded. Thread-safe — the internal Mutex
    /// serialises against concurrent transcribe_sync calls.
    pub fn unload(&self) {
        let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
        *guard = None;
        let mut id_guard = self
            .loaded_model_id
            .lock()
            .unwrap_or_else(|e| e.into_inner());
        *id_guard = None;
    }

    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()
    }

    /// Capabilities of the currently-loaded backend. Returns `None`
    /// when nothing is loaded. Callers (live capture WAV writer, #19)
    /// read sample_rate from here.
    pub fn capabilities(&self) -> Option<TranscriberCapabilities> {
        let guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
        guard.as_ref().map(|b| b.capabilities())
    }

    /// 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 = backend.transcribe_sync(audio.samples(), options)?;
        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<Box<dyn Transcriber + Send>> {
    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(Box::new(SpeechModelAdapter(Box::new(
        ParakeetWordGranularity(model),
    ))))
}

/// Load a Whisper model from a GGML file path via whisper-rs.
#[cfg(feature = "whisper")]
pub fn load_whisper(model_path: &Path) -> Result<Box<dyn Transcriber + Send>> {
    let backend = WhisperRsBackend::load(model_path)
        .map_err(|e| KonError::TranscriptionFailed(format!("Failed to load Whisper: {e}")))?;
    Ok(Box::new(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());
        assert!(engine.capabilities().is_none());
    }
}