Lumotia/crates/transcription/src/orchestrator.rs

//! `Orchestrator` is the single entry point for transcription. Tauri
//! commands resolve a provider through it; nothing else calls a
//! provider directly. This is where the AGPL OEM trait boundary meets
//! Lumotia's local engines.
//!
//! `LocalProviderAdapter` lives here, not as an `impl
//! TranscriptionProvider for LocalEngine`. The adapter wraps an
//! `Arc<LocalEngine>` and bridges the synchronous `Transcriber` trait
//! to the async `TranscriptionProvider` interface via
//! `tokio::task::spawn_blocking`. Keeping the adapter in the
//! orchestrator (rather than in `local_engine.rs`) means the
//! transcription crate's core types do not need to depend on
//! `async_trait`, and the dependency edge stays one-directional:
//! `cloud_providers` defines the trait, `transcription` implements it
//! via adapter without leaking async-runtime requirements onto
//! `Transcriber`.

use std::sync::Arc;

use async_trait::async_trait;
use lumotia_cloud_providers::{
    CostClass, EngineProfile, ProviderCapabilities, ProviderId, ProviderKind, ProviderTranscript,
    TranscriptionProvider,
};
use lumotia_core::error::{Error, Result};
use lumotia_core::types::{AudioSamples, TranscriptionOptions};

use crate::local_engine::LocalEngine;
use crate::registry::EngineRegistry;

/// Wraps a `LocalEngine` and presents the async `TranscriptionProvider`
/// interface upward. One adapter per local engine instance; multiple
/// engines (Whisper, Parakeet) live as multiple adapters in the
/// registry.
pub struct LocalProviderAdapter {
    provider_id: ProviderId,
    engine: Arc<LocalEngine>,
}

impl LocalProviderAdapter {
    pub fn new(provider_id: ProviderId, engine: Arc<LocalEngine>) -> Self {
        Self {
            provider_id,
            engine,
        }
    }

    /// Direct access to the wrapped engine. Used by warmup, model
    /// management, and the GPU-sequencing guard, none of which want to
    /// route through the async trait surface.
    pub fn engine(&self) -> Arc<LocalEngine> {
        self.engine.clone()
    }
}

#[async_trait]
impl TranscriptionProvider for LocalProviderAdapter {
    fn provider_id(&self) -> ProviderId {
        self.provider_id.clone()
    }

    fn kind(&self) -> ProviderKind {
        ProviderKind::Local
    }

    fn capabilities(&self) -> ProviderCapabilities {
        let local = self.engine.capabilities();
        ProviderCapabilities {
            sample_rate: local
                .map(|c| c.sample_rate)
                .unwrap_or(lumotia_core::constants::WHISPER_SAMPLE_RATE),
            channels: local.map(|c| c.channels).unwrap_or(1),
            initial_prompt_supported: local.map(|c| c.supports_initial_prompt).unwrap_or(false),
            language_hint_supported: true,
            streaming_supported: false,
            cost_class: CostClass::Free,
        }
    }

    async fn prepare(&self, _profile: &EngineProfile) -> Result<()> {
        if self.engine.is_loaded() {
            Ok(())
        } else {
            Err(Error::EngineNotLoaded)
        }
    }

    async fn transcribe(
        &self,
        audio: AudioSamples,
        options: TranscriptionOptions,
    ) -> Result<ProviderTranscript> {
        let engine = self.engine.clone();
        let timed = tokio::task::spawn_blocking(move || engine.transcribe_sync(&audio, &options))
            .await
            .map_err(|e| Error::TranscriptionFailed(format!("Task join error: {e}")))??;
        Ok(ProviderTranscript {
            transcript: timed.transcript,
            inference_ms: timed.inference_ms,
        })
    }
}

/// The single entry point through which all transcription flows.
/// Resolves a provider against the registry, runs `prepare` if the
/// caller has not already done so, and dispatches `transcribe`.
pub struct Orchestrator {
    registry: Arc<EngineRegistry>,
}

impl Orchestrator {
    pub fn new(registry: Arc<EngineRegistry>) -> Self {
        Self { registry }
    }

    /// Underlying registry. Exposed so the UI can list providers and
    /// query capabilities without going through a transcribe call.
    pub fn registry(&self) -> Arc<EngineRegistry> {
        self.registry.clone()
    }

    /// Resolve the provider for a profile. Returns a clear error when
    /// the profile names an unregistered engine.
    pub fn resolve(&self, profile: &EngineProfile) -> Result<Arc<dyn TranscriptionProvider>> {
        self.registry
            .get(&profile.engine_id)
            .ok_or_else(|| Error::ProviderNotRegistered(profile.engine_id.to_string()))
    }

    /// Transcribe audio using the provider named in the profile. The
    /// orchestrator builds `TranscriptionOptions` from the profile and
    /// delegates to the provider.
    ///
    /// Phase A scope: this is the new entry point. Existing call sites
    /// (`commands/transcription.rs::transcribe_file`) still call
    /// `LocalEngine` directly; their rewire is a follow-up commit so
    /// the chunking + post-processing logic moves cleanly without
    /// inflating this diff. See `KNOWN-ISSUES.md` KI-06.
    pub async fn transcribe(
        &self,
        audio: AudioSamples,
        profile: &EngineProfile,
    ) -> Result<ProviderTranscript> {
        let provider = self.resolve(profile)?;
        let options = profile.to_options();
        provider.transcribe(audio, options).await
    }
}

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

    use std::sync::atomic::{AtomicUsize, Ordering};

    use lumotia_core::types::{Segment, Transcript};

    /// Mock provider that returns a canned transcript and counts
    /// invocations. Used to validate the orchestrator's dispatch logic
    /// without booting a real model.
    struct CannedProvider {
        id: ProviderId,
        calls: Arc<AtomicUsize>,
        canned_text: String,
    }

    #[async_trait]
    impl TranscriptionProvider for CannedProvider {
        fn provider_id(&self) -> ProviderId {
            self.id.clone()
        }
        fn kind(&self) -> ProviderKind {
            ProviderKind::Local
        }
        fn capabilities(&self) -> ProviderCapabilities {
            ProviderCapabilities {
                sample_rate: 16_000,
                channels: 1,
                initial_prompt_supported: true,
                language_hint_supported: true,
                streaming_supported: false,
                cost_class: CostClass::Free,
            }
        }
        async fn prepare(&self, _profile: &EngineProfile) -> Result<()> {
            Ok(())
        }
        async fn transcribe(
            &self,
            audio: AudioSamples,
            _options: TranscriptionOptions,
        ) -> Result<ProviderTranscript> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            Ok(ProviderTranscript {
                transcript: Transcript::new(
                    vec![Segment {
                        start: 0.0,
                        end: audio.duration_secs(),
                        text: self.canned_text.clone(),
                    }],
                    "en".to_string(),
                    audio.duration_secs(),
                ),
                inference_ms: 7,
            })
        }
    }

    fn build_registry_with(id: &str, text: &str, calls: Arc<AtomicUsize>) -> Arc<EngineRegistry> {
        let mut registry = EngineRegistry::new(ProviderId::new(id));
        registry.register(Arc::new(CannedProvider {
            id: ProviderId::new(id),
            calls,
            canned_text: text.to_string(),
        }));
        Arc::new(registry)
    }

    fn one_second_of_silence() -> AudioSamples {
        AudioSamples::mono_16khz(vec![0.0_f32; 16_000])
    }

    #[tokio::test]
    async fn orchestrator_dispatches_to_named_provider() {
        let calls = Arc::new(AtomicUsize::new(0));
        let registry = build_registry_with("canned", "hello lumotia", calls.clone());
        let orchestrator = Orchestrator::new(registry);

        let profile = EngineProfile::new(ProviderId::new("canned"));
        let result = orchestrator
            .transcribe(one_second_of_silence(), &profile)
            .await
            .expect("dispatch succeeds");

        assert_eq!(result.transcript.text(), "hello lumotia");
        assert_eq!(result.inference_ms, 7);
        assert_eq!(calls.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn orchestrator_returns_clear_error_for_unregistered_provider() {
        let calls = Arc::new(AtomicUsize::new(0));
        let registry = build_registry_with("canned", "_", calls);
        let orchestrator = Orchestrator::new(registry);

        let profile = EngineProfile::new(ProviderId::new("not-registered"));
        let err = orchestrator
            .transcribe(one_second_of_silence(), &profile)
            .await
            .expect_err("unregistered provider must error");

        let msg = err.to_string();
        assert!(
            msg.contains("not-registered"),
            "error must name the missing provider, got: {msg}"
        );
    }

    #[tokio::test]
    async fn orchestrator_routes_initial_prompt_through_options() {
        let calls = Arc::new(AtomicUsize::new(0));
        let registry = build_registry_with("canned", "transcript", calls.clone());
        let orchestrator = Orchestrator::new(registry);

        let mut profile = EngineProfile::new(ProviderId::new("canned"));
        profile.language = Some("en".to_string());
        profile.initial_prompt = Some("Lumotia".to_string());

        let _ = orchestrator
            .transcribe(one_second_of_silence(), &profile)
            .await
            .expect("dispatch succeeds");

        assert_eq!(calls.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn local_provider_adapter_is_object_safe() {
        let _: Option<Arc<dyn TranscriptionProvider>> = None;
    }
}