local_transcriber.rs

use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Instant;

use color_eyre::eyre::{Result, WrapErr, eyre};
use transcribe_cpp::{Backend, Model, ModelOptions, RunOptions, TimestampKind};

use crate::gguf_session::OfflineGgufSession;
use crate::transcription_models::{
    TranscriptionSelection, download_with_progress, model_path, validate,
};

/// A word or segment's start and end (ms from the clip's start) and text.
pub type TimedSpan = (i64, i64, String);

pub struct LocalTranscriber {
    session: OfflineGgufSession,
    selection: TranscriptionSelection,
    options: RunOptions,
    device_label: String,
    max_audio_samples: Option<usize>,
    stream: Option<transcribe_cpp::StreamOptions>,
    /// The finest timestamps the model gives, for a capture's or a
    /// meeting's word times.
    finest_timestamps: TimestampKind,
}

pub fn install(selection: &TranscriptionSelection, canceled: &AtomicBool) -> Result<()> {
    let model = validate(selection)?;
    let downloaded = AtomicU64::new(0);
    println!("Installing {} ({})...", model.name, model.size_label());
    let path = download_with_progress(model, canceled, &downloaded)?;
    println!(
        "Installed {} bytes at {}",
        downloaded.load(Ordering::Relaxed),
        path.display()
    );
    Ok(())
}

pub fn devices() -> Result<Vec<String>> {
    transcribe_cpp::init_logging();
    transcribe_cpp::init_backends_default()
        .wrap_err("could not initialize transcription backends")?;
    Ok(transcribe_cpp::devices()
        .into_iter()
        .map(|device| {
            format!(
                "{}\t{}\t{} MB\t{}",
                device
                    .index
                    .map_or_else(|| "-".into(), |index| index.to_string()),
                device.kind,
                device.memory_total / (1024 * 1024),
                if device.description.is_empty() {
                    device.name
                } else {
                    device.description
                }
            )
        })
        .collect())
}

/// Initializes transcribe.cpp's backends, which creates its Vulkan instance,
/// on the calling thread. The Linux process calls this before its listener
/// or window can start: two Vulkan instances created at once in one process
/// (ggml's here, wgpu's for the window) crash inside the Vulkan loader.
/// Later model loads reuse the instance.
#[cfg(target_os = "linux")]
pub fn initialize_backends() {
    transcribe_cpp::init_logging();
    match transcribe_cpp::init_backends_default() {
        // Enumerating the devices makes sure the Vulkan backend exists now.
        Ok(()) => drop(transcribe_cpp::devices()),
        Err(error) => tracing::warn!(%error, "could not initialize transcription backends"),
    }
}

pub fn transcribe_wav_with_selection(
    path: &Path,
    selection: &TranscriptionSelection,
) -> Result<String> {
    let samples = read_wav_for_parakeet(path)?;
    LocalTranscriber::load(selection)?.transcribe(&samples)
}

pub fn read_wav_for_parakeet(path: &Path) -> Result<Vec<f32>> {
    let mut reader = hound::WavReader::open(path)
        .wrap_err_with(|| format!("could not open {}", path.display()))?;
    let spec = reader.spec();
    let samples = match spec.sample_format {
        hound::SampleFormat::Float => reader
            .samples::<f32>()
            .collect::<std::result::Result<Vec<_>, _>>()?,
        hound::SampleFormat::Int if spec.bits_per_sample <= 16 => reader
            .samples::<i16>()
            .map(|sample| sample.map(|sample| sample as f32 / i16::MAX as f32))
            .collect::<std::result::Result<Vec<_>, _>>()?,
        hound::SampleFormat::Int => reader
            .samples::<i32>()
            .map(|sample| sample.map(|sample| sample as f32 / i32::MAX as f32))
            .collect::<std::result::Result<Vec<_>, _>>()?,
    };
    let channels = usize::from(spec.channels);
    let mono = samples
        .chunks(channels)
        .map(|frame| frame.iter().sum::<f32>() / channels as f32)
        .collect::<Vec<_>>();
    let mut samples = crate::dictation::resample_for_parakeet(&mono, spec.sample_rate);
    crate::dictation::pad_for_parakeet(&mut samples);
    Ok(samples)
}

impl LocalTranscriber {
    pub fn load(selection: &TranscriptionSelection) -> Result<Self> {
        transcribe_cpp::init_logging();
        transcribe_cpp::init_backends_default()
            .wrap_err("could not initialize transcription backends")?;
        let definition = validate(selection)?;
        if !crate::transcription_models::is_installed(definition, &selection.language) {
            return Err(eyre!(
                "{} is not installed; run `cargo run -- model install`",
                definition.name
            ));
        }
        let path = model_path(definition)?;
        let model = Model::load_with(
            &path,
            &ModelOptions {
                backend: Backend::Auto,
                ..Default::default()
            },
        )
        .wrap_err_with(|| format!("could not load transcription model from {}", path.display()))?;
        let device = model.device()?;
        let variant = model.variant();
        let architecture = model.arch();
        let capabilities = model.capabilities()?;
        let language = crate::gguf_session::validate_gguf_artifact(
            &path,
            definition,
            &architecture,
            &variant,
            selection,
            capabilities.supports_language_detect,
        )?;
        let mut options = RunOptions {
            timestamps: if capabilities.max_timestamp_kind == TimestampKind::None {
                TimestampKind::None
            } else {
                TimestampKind::Segment
            },
            language,
            ..Default::default()
        };
        crate::gguf_session::apply_recognition_hints(
            &mut options,
            &model,
            definition,
            &selection.recognition_hints,
        )?;
        let device_label = format!("{} ({})", device.description, device.kind);
        tracing::info!(
            backend = device.kind,
            device = device.description,
            variant,
            "loaded local transcription model"
        );
        let session = OfflineGgufSession::new(model)?;
        let mut transcriber = Self {
            session,
            selection: selection.clone(),
            options,
            device_label,
            max_audio_samples: crate::transcription_models::max_audio_chunk_samples(
                &architecture,
                capabilities.max_audio_ms,
            ),
            stream: None,
            finest_timestamps: match capabilities.max_timestamp_kind {
                TimestampKind::Word | TimestampKind::Token => TimestampKind::Word,
                kind => kind,
            },
        };
        let started = Instant::now();
        transcriber.transcribe(&vec![0.0; 24_000])?;
        tracing::info!(
            prewarm_ms = started.elapsed().as_millis(),
            "prewarmed local transcription model"
        );
        let stream_run = transcriber.stream_run_options();
        transcriber.stream = transcriber.session.probe_stream(&stream_run, &architecture);
        Ok(transcriber)
    }

    /// Streams decode exactly as offline, without the recognition hints'
    /// prompt (Whisper's run-slot extension or the context prompt).
    fn stream_run_options(&self) -> RunOptions {
        RunOptions {
            family: None,
            prompt: None,
            ..self.options.clone()
        }
    }

    pub(crate) fn stream_options(&self) -> Option<&transcribe_cpp::StreamOptions> {
        self.stream.as_ref()
    }

    pub(crate) fn with_stream<R>(
        &mut self,
        stream_options: &transcribe_cpp::StreamOptions,
        cancel: &transcribe_cpp::CancelToken,
        drive: impl FnOnce(&mut transcribe_cpp::Stream<'_>) -> R,
    ) -> transcribe_cpp::Result<R> {
        let options = self.stream_run_options();
        self.session
            .with_stream(&options, stream_options, cancel, drive)
    }

    /// The selection this model was loaded for, so the transcription engine
    /// can say which model it holds.
    pub(crate) fn gguf_selection(&self) -> Option<&TranscriptionSelection> {
        Some(&self.selection)
    }

    pub(crate) fn matches_selection(&self, selection: &TranscriptionSelection) -> bool {
        &self.selection == selection
    }

    /// Brings an idle model's weights back before audio arrives, as loading
    /// does; a model that ran recently is already warm.
    pub(crate) fn prewarm(&mut self) -> Result<()> {
        if self.session.ran_recently() {
            return Ok(());
        }
        self.transcribe(&vec![0.0; 24_000])
            .wrap_err("could not prewarm transcription model")
            .map(drop)
    }

    /// Transcribes a complete 16 kHz clip, padded as a dictation's clip is.
    pub(crate) fn transcribe_clip(&mut self, mut samples: Vec<f32>) -> Result<String> {
        crate::dictation::pad_for_parakeet(&mut samples);
        self.transcribe(&samples)
    }

    pub fn device_label(&self) -> &str {
        &self.device_label
    }

    /// Text and word times (ms from the clip's start) for a capture, in
    /// chunks the model takes whole. A model without word times gives its
    /// segments instead.
    pub fn transcribe_words(&mut self, samples: &[f32]) -> Result<(String, Vec<TimedSpan>)> {
        let transcript = self.timed(samples)?;
        let spans = if transcript.words.is_empty() {
            transcript.segments
        } else {
            transcript.words
        };
        let words = spans
            .into_iter()
            .map(|word| (word.start_ms, word.end_ms, word.text))
            .collect();
        Ok((transcript.text, words))
    }

    /// One meeting chunk with segment and word timings, padded as dictation
    /// pads a clip.
    pub(crate) fn transcribe_segments(
        &mut self,
        mut samples: Vec<f32>,
    ) -> Result<crate::transcript::Transcript> {
        crate::dictation::pad_for_parakeet(&mut samples);
        self.timed(&samples)
    }

    /// Text, segments, and words, in chunks the model takes whole.
    fn timed(&mut self, samples: &[f32]) -> Result<crate::transcript::Transcript> {
        use crate::transcript::TranscriptSegment;
        const CHUNK: usize = 30 * 16_000;
        let chunk = self.max_audio_samples.map_or(CHUNK, |max| max.min(CHUNK));
        let options = RunOptions {
            timestamps: self.finest_timestamps,
            ..self.options.clone()
        };
        let span = |offset: i64, t0: i64, t1: i64, text: String| {
            let text = text.trim().to_owned();
            (!text.is_empty()).then(|| TranscriptSegment {
                start_ms: t0.max(0) + offset,
                end_ms: t1.max(t0) + offset,
                text,
            })
        };
        let mut timed = crate::transcript::Transcript::default();
        let mut texts = Vec::new();
        for (index, part) in samples.chunks(chunk.max(1)).enumerate() {
            let offset = (index * chunk / 16) as i64;
            let transcript = self
                .session
                .run(part, &options)
                .wrap_err("transcription failed")?;
            let text = transcript.text.trim();
            if !text.is_empty() {
                texts.push(text.to_owned());
            }
            timed.segments.extend(
                transcript
                    .segments
                    .into_iter()
                    .filter_map(|s| span(offset, s.t0_ms, s.t1_ms, s.text)),
            );
            timed.words.extend(
                transcript
                    .words
                    .into_iter()
                    .filter_map(|w| span(offset, w.t0_ms, w.t1_ms, w.text)),
            );
        }
        timed.text = texts.join(" ");
        Ok(timed)
    }

    pub fn transcribe(&mut self, samples: &[f32]) -> Result<String> {
        self.session
            .transcribe_chunked(samples, &self.options, self.max_audio_samples)
    }
}

/// The local API's own model, loaded for the selection a host asks for
/// when the dictation engine holds another one.
#[derive(Default)]
pub(crate) struct WarmLocal {
    active: Option<LocalTranscriber>,
}

impl WarmLocal {
    /// The loaded model, without switching or loading one.
    pub(crate) fn active(&mut self) -> Option<&mut LocalTranscriber> {
        self.active.as_mut()
    }

    pub(crate) fn activate(
        &mut self,
        selection: &TranscriptionSelection,
    ) -> Result<&mut LocalTranscriber> {
        if !self
            .active
            .as_ref()
            .is_some_and(|model| model.matches_selection(selection))
        {
            self.active = None;
            self.active = Some(LocalTranscriber::load(selection)?);
        }
        Ok(self
            .active
            .as_mut()
            .expect("activated transcriber must be available"))
    }
}

/// A running listener's model: the one it transcribes with, and the one it
/// changes to. A new selection loads on a worker thread while dictation goes
/// on with the old model; the listener swaps it in once nothing records or
/// waits on the old one. A model the host already prepared skips the load.
pub(crate) struct ModelSwap {
    loaded: TranscriptionSelection,
    loading: Option<(
        TranscriptionSelection,
        std::sync::mpsc::Receiver<std::result::Result<LocalTranscriber, String>>,
    )>,
    ready: Option<(TranscriptionSelection, LocalTranscriber)>,
}

impl ModelSwap {
    pub(crate) fn new(loaded: &TranscriptionSelection) -> Self {
        Self {
            loaded: runnable(loaded),
            loading: None,
            ready: None,
        }
    }

    /// The selection the listener will transcribe with once swaps settle.
    fn target(&self) -> &TranscriptionSelection {
        self.ready
            .as_ref()
            .map(|(selection, _)| selection)
            .or(self.loading.as_ref().map(|(selection, _)| selection))
            .unwrap_or(&self.loaded)
    }

    /// A model the host loaded for `selection` (Linux prepares models
    /// beside the listener; Windows lets the listener load them).
    #[cfg_attr(target_os = "windows", allow(dead_code))]
    pub(crate) fn offer(
        &mut self,
        selection: &TranscriptionSelection,
        transcriber: LocalTranscriber,
    ) {
        self.loading = None;
        self.ready = Some((runnable(selection), transcriber));
    }

    /// Settings now select `selection`; loads it unless it is already the
    /// target.
    pub(crate) fn want(&mut self, selection: &TranscriptionSelection) {
        let selection = runnable(selection);
        if &selection == self.target() {
            return;
        }
        self.ready = None;
        if selection == self.loaded {
            self.loading = None;
            return;
        }
        let (sender, receiver) = std::sync::mpsc::channel();
        let load = selection.clone();
        let spawned = std::thread::Builder::new()
            .name("hex-model-swap".into())
            .spawn(move || {
                let _ = sender
                    .send(LocalTranscriber::load(&load).map_err(|error| format!("{error:#}")));
            });
        match spawned {
            Ok(_) => self.loading = Some((selection, receiver)),
            Err(error) => tracing::warn!(%error, "could not load the new transcription model"),
        }
    }

    /// Collects a finished load; a failure keeps the current model.
    pub(crate) fn poll(&mut self) -> Option<SwapFailed> {
        let (_, receiver) = self.loading.as_ref()?;
        let result = match receiver.try_recv() {
            Ok(result) => result,
            Err(std::sync::mpsc::TryRecvError::Empty) => return None,
            Err(std::sync::mpsc::TryRecvError::Disconnected) => {
                Err("the model loader stopped unexpectedly".into())
            }
        };
        let (selection, _) = self.loading.take()?;
        match result {
            Ok(transcriber) => {
                self.ready = Some((selection, transcriber));
                None
            }
            Err(error) => Some(SwapFailed {
                wanted: selection,
                kept: self.loaded.clone(),
                error: format!("could not switch the transcription model: {error}"),
            }),
        }
    }

    /// The model to swap in now, if one is ready.
    pub(crate) fn take_ready(&mut self) -> Option<LocalTranscriber> {
        let (selection, transcriber) = self.ready.take()?;
        self.loaded = selection;
        Some(transcriber)
    }
}

/// A model change that failed to load; the listener still runs `kept`.
/// Windows hands it to the host, which puts settings back on `kept`; Linux
/// only reports the error.
#[derive(Debug)]
#[cfg_attr(target_os = "linux", allow(dead_code))]
pub(crate) struct SwapFailed {
    pub(crate) wanted: TranscriptionSelection,
    pub(crate) kept: TranscriptionSelection,
    pub(crate) error: String,
}

/// `selection` without recognition hints a model cannot use, which
/// `validate` would reject.
fn runnable(selection: &TranscriptionSelection) -> TranscriptionSelection {
    let mut selection = selection.clone();
    if !crate::transcription_models::definition(selection.model).supports_recognition_hints {
        selection.recognition_hints = String::new();
    }
    selection
}

#[cfg(test)]
mod model_swap_tests {
    use super::*;
    use crate::transcription_models::TranscriptionModelId;

    #[test]
    fn an_unchanged_or_hint_only_change_on_a_hintless_model_loads_nothing() {
        let selection = TranscriptionSelection {
            model: TranscriptionModelId::default(),
            ..Default::default()
        };
        let mut swap = ModelSwap::new(&selection);
        swap.want(&selection);
        let mut hinted = selection.clone();
        hinted.recognition_hints = "Hex".into();
        if !crate::transcription_models::definition(hinted.model).supports_recognition_hints {
            swap.want(&hinted);
        }
        assert!(swap.loading.is_none());
        assert!(swap.take_ready().is_none());
        assert!(swap.poll().is_none());
    }
}