Lumotia/crates/transcription/src/local_engine.rs

use std::path::Path;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex};
use std::time::Instant;

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

use lumotia_core::error::{Error, Result};
use lumotia_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: lumotia_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| Error::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())
    }

    /// SAFETY: `transcribe-rs` owns the Parakeet decoder behind a `&mut
    /// self` call that does not surface a cancellation hook. The best we
    /// can do without forking transcribe-rs is short-circuit BEFORE the
    /// decode call when the live session has already requested abort —
    /// the same boundary the `Drop for InferenceTask` cancellation
    /// route arrives through. Once the decode is in flight it runs to
    /// completion, but the live session's `drain_inference` timeout
    /// still drops our receiver, so the wedged orphan thread exits the
    /// instant it tries to send its result. The engine `Mutex` is held
    /// only across THIS call, so a future task will not deadlock — it
    /// simply queues until the orphan releases. Documenting the
    /// uncancellable middle so future audits don't get surprised.
    fn transcribe_sync_with_abort(
        &mut self,
        samples: &[f32],
        options: &TranscriptionOptions,
        abort_flag: Arc<AtomicBool>,
    ) -> Result<Vec<Segment>> {
        if abort_flag.load(std::sync::atomic::Ordering::Relaxed) {
            return Err(Error::TranscriptionFailed(
                "transcription aborted before decoder dispatch".to_string(),
            ));
        }
        self.transcribe_sync(samples, options)
    }
}

/// 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(Error::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,
        })
    }

    /// Cancellable variant of `transcribe_sync`. Pipes `abort_flag`
    /// through to the backend so the live-session drain timeout can
    /// break a wedged whisper-rs decode out of its loop. Every backend
    /// MUST implement the trait method (no default impl) — Parakeet's
    /// adapter, for example, honours the flag at the pre-decode
    /// boundary even though the in-flight decode itself is opaque.
    pub fn transcribe_sync_with_abort(
        &self,
        audio: &AudioSamples,
        options: &TranscriptionOptions,
        abort_flag: Arc<AtomicBool>,
    ) -> Result<TimedTranscript> {
        let mut guard = self.engine.lock().unwrap_or_else(|e| e.into_inner());
        let backend = guard.as_mut().ok_or(Error::EngineNotLoaded)?;

        let start = Instant::now();
        let segments = backend.transcribe_sync_with_abort(audio.samples(), options, abort_flag)?;
        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 Lumotia 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| Error::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| Error::TranscriptionFailed(format!("Failed to load Whisper: {e}")))?;
    Ok(Box::new(backend))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::Ordering;
    use transcribe_rs::{
        ModelCapabilities, TranscribeError, TranscribeOptions, TranscriptionResult,
    };

    #[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());
    }

    /// Minimal fake `SpeechModel` for the Parakeet adapter tests. Records
    /// whether `transcribe_raw` was called so we can prove the pre-decode
    /// abort short-circuit actually fires.
    struct FakeSpeechModel {
        call_count: Arc<std::sync::atomic::AtomicUsize>,
    }

    impl transcribe_rs::SpeechModel for FakeSpeechModel {
        fn capabilities(&self) -> ModelCapabilities {
            ModelCapabilities {
                name: "fake",
                engine_id: "fake",
                sample_rate: 16_000,
                languages: &[],
                supports_timestamps: false,
                supports_translation: false,
                supports_streaming: false,
            }
        }

        fn transcribe_raw(
            &mut self,
            _samples: &[f32],
            _options: &TranscribeOptions,
        ) -> std::result::Result<TranscriptionResult, TranscribeError> {
            self.call_count
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            Ok(TranscriptionResult {
                text: String::new(),
                segments: Some(Vec::new()),
            })
        }
    }

    #[test]
    fn speech_model_adapter_short_circuits_when_abort_set_pre_dispatch() {
        // Lifecycle-2 regression: without an explicit
        // `transcribe_sync_with_abort` impl, SpeechModelAdapter used to
        // inherit a default that silently dropped the abort flag and
        // ran the decoder anyway. With the trait method made required,
        // the adapter now checks the flag at the safest available
        // boundary (pre-decode) and short-circuits.
        let call_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let model = Box::new(FakeSpeechModel {
            call_count: call_count.clone(),
        });
        let mut adapter = SpeechModelAdapter(model);

        let abort = Arc::new(AtomicBool::new(true));
        let options = TranscriptionOptions::default();
        let res = adapter.transcribe_sync_with_abort(&[0.0_f32; 16], &options, abort);
        assert!(res.is_err(), "pre-set abort flag must short-circuit");
        assert_eq!(
            call_count.load(Ordering::Relaxed),
            0,
            "underlying decoder must NOT be called when abort was set before dispatch"
        );
    }

    #[test]
    fn speech_model_adapter_dispatches_decoder_when_abort_clear() {
        // Companion to the short-circuit test: when the abort flag is
        // clear at dispatch time, the adapter must still call the
        // underlying decoder. Without this we'd have killed
        // transcription entirely.
        let call_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
        let model = Box::new(FakeSpeechModel {
            call_count: call_count.clone(),
        });
        let mut adapter = SpeechModelAdapter(model);

        let abort = Arc::new(AtomicBool::new(false));
        let options = TranscriptionOptions::default();
        let res = adapter.transcribe_sync_with_abort(&[0.0_f32; 16], &options, abort);
        assert!(res.is_ok(), "clear abort flag must allow dispatch");
        assert_eq!(
            call_count.load(Ordering::Relaxed),
            1,
            "decoder must run exactly once when abort flag is clear"
        );
    }
}