Lumotia/src-tauri/src/commands/transcription.rs

// Tauri command handlers must match the frontend's invoke() parameter lists,
// so the argument counts are dictated by the Svelte code.
#![allow(clippy::too_many_arguments)]

use std::path::Path;
use std::sync::Arc;

use tauri::Emitter;

use crate::commands::build_initial_prompt;
use crate::commands::models::{default_model_id_for_engine, ensure_model_loaded};
use crate::AppState;
use kon_ai_formatting::{post_process_segments, FormatMode, PostProcessOptions};
use kon_core::constants::WHISPER_SAMPLE_RATE;
use kon_core::types::{AudioSamples, Segment, Transcript, TranscriptionOptions};

const PARAKEET_CHUNK_THRESHOLD_SECS: usize = 18;
const PARAKEET_CHUNK_SECS: usize = 15;
const PARAKEET_CHUNK_OVERLAP_SECS: usize = 1;
const FILE_CHUNK_THRESHOLD_SECS: usize = 8 * 60;
const FILE_CHUNK_SECS: usize = 3 * 60;
const FILE_CHUNK_OVERLAP_SECS: usize = 2;

struct ChunkingStrategy {
    chunk_samples: usize,
    overlap_samples: usize,
}

fn pick_engine(
    state: &AppState,
    engine: &str,
) -> Result<Arc<kon_transcription::LocalEngine>, String> {
    match engine {
        "whisper" => Ok(state.whisper_engine.clone()),
        "parakeet" => Ok(state.parakeet_engine.clone()),
        other => Err(format!("Unknown engine: {other}")),
    }
}

fn pick_chunking_strategy(engine_name: &str, sample_count: usize) -> Option<ChunkingStrategy> {
    let samples_per_second = WHISPER_SAMPLE_RATE as usize;
    match engine_name {
        "parakeet" if sample_count > PARAKEET_CHUNK_THRESHOLD_SECS * samples_per_second => {
            Some(ChunkingStrategy {
                chunk_samples: PARAKEET_CHUNK_SECS * samples_per_second,
                overlap_samples: PARAKEET_CHUNK_OVERLAP_SECS * samples_per_second,
            })
        }
        _ if sample_count > FILE_CHUNK_THRESHOLD_SECS * samples_per_second => {
            Some(ChunkingStrategy {
                chunk_samples: FILE_CHUNK_SECS * samples_per_second,
                overlap_samples: FILE_CHUNK_OVERLAP_SECS * samples_per_second,
            })
        }
        _ => None,
    }
}

fn trim_overlap_segments(segments: &mut Vec<Segment>, trim_before_secs: f64) {
    if trim_before_secs <= 0.0 {
        return;
    }

    segments.retain(|segment| segment.end > trim_before_secs);
    for segment in segments.iter_mut() {
        if segment.start < trim_before_secs {
            segment.start = trim_before_secs;
        }
    }
}

fn transcribe_samples_sync(
    engine: Arc<kon_transcription::LocalEngine>,
    engine_name: &str,
    samples: Vec<f32>,
    options: TranscriptionOptions,
) -> Result<kon_transcription::TimedTranscript, String> {
    let Some(strategy) = pick_chunking_strategy(engine_name, samples.len()) else {
        let audio = AudioSamples::mono_16khz(samples);
        return engine
            .transcribe_sync(&audio, &options)
            .map_err(|e| e.to_string());
    };

    let total_duration_secs = samples.len() as f64 / WHISPER_SAMPLE_RATE as f64;
    let stride = strategy
        .chunk_samples
        .saturating_sub(strategy.overlap_samples)
        .max(1);
    let chunk_count = ((samples.len().saturating_sub(1)) / stride) + 1;
    eprintln!(
        "[transcription] chunking {total_duration_secs:.2}s of {engine_name} audio into {chunk_count} chunk(s)"
    );

    let mut all_segments = Vec::new();
    let mut total_inference_ms = 0u64;
    let mut chunk_start = 0usize;

    while chunk_start < samples.len() {
        let chunk_end = (chunk_start + strategy.chunk_samples).min(samples.len());
        let chunk_audio = AudioSamples::mono_16khz(samples[chunk_start..chunk_end].to_vec());
        let timed = engine
            .transcribe_sync(&chunk_audio, &options)
            .map_err(|e| e.to_string())?;
        total_inference_ms = total_inference_ms.saturating_add(timed.inference_ms);

        let mut chunk_segments = timed.transcript.segments().to_vec();
        if chunk_start > 0 {
            trim_overlap_segments(
                &mut chunk_segments,
                strategy.overlap_samples as f64 / WHISPER_SAMPLE_RATE as f64,
            );
        }

        let chunk_offset_secs = chunk_start as f64 / WHISPER_SAMPLE_RATE as f64;
        for segment in &mut chunk_segments {
            segment.start += chunk_offset_secs;
            segment.end += chunk_offset_secs;
        }
        all_segments.extend(chunk_segments);

        if chunk_end >= samples.len() {
            break;
        }
        chunk_start = chunk_end.saturating_sub(strategy.overlap_samples);
    }

    Ok(kon_transcription::TimedTranscript {
        transcript: Transcript::new(
            all_segments,
            options.language.clone().unwrap_or_else(|| "en".to_string()),
            total_duration_secs,
        ),
        inference_ms: total_inference_ms,
    })
}

/// Transcribe raw PCM f32 samples (Whisper). Emits "transcription-result" event.
#[tauri::command]
pub async fn transcribe_pcm(
    state: tauri::State<'_, AppState>,
    app: tauri::AppHandle,
    samples: Vec<f32>,
    chunk_id: u32,
    language: String,
    initial_prompt: String,
    remove_fillers: bool,
    british_english: bool,
    anti_hallucination: bool,
    format_mode: String,
    profile_id: Option<String>,
) -> Result<(), String> {
    let resolved_profile_id =
        profile_id.unwrap_or_else(|| kon_storage::DEFAULT_PROFILE_ID.to_string());

    let profile = kon_storage::database::get_profile(&state.db, &resolved_profile_id)
        .await
        .map_err(|e| e.to_string())?
        .ok_or_else(|| format!("Profile {resolved_profile_id} not found"))?;

    let profile_terms: Vec<String> =
        kon_storage::database::list_profile_terms(&state.db, &resolved_profile_id)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .map(|t| t.term)
            .collect();

    let engine = state.whisper_engine.clone();
    let options = TranscriptionOptions {
        language: Some(language),
        initial_prompt: build_initial_prompt(
            &initial_prompt,
            &profile.initial_prompt,
            &profile_terms,
        ),
    };

    let timed = tokio::task::spawn_blocking(move || {
        let audio = AudioSamples::mono_16khz(samples);
        engine
            .transcribe_sync(&audio, &options)
            .map_err(|e| e.to_string())
    })
    .await
    .map_err(|e| e.to_string())??;

    let dictionary_terms = profile_terms.clone();

    let mut segments: Vec<Segment> = timed.transcript.segments().to_vec();
    let raw_text = join_segment_text(&segments);
    post_process_segments(
        &mut segments,
        &PostProcessOptions {
            remove_fillers,
            british_english,
            anti_hallucination,
            format_mode: FormatMode::parse(&format_mode),
            dictionary_terms,
        },
        Some(state.llm_engine.as_ref()),
    );

    app.emit(
        "transcription-result",
        serde_json::json!({
            "status": "transcription",
            "segments": segments,
            "language": timed.transcript.language(),
            "duration": timed.transcript.duration(),
            "chunk_id": chunk_id,
            "inference_ms": timed.inference_ms,
            "raw_text": raw_text,
        }),
    )
    .map_err(|e| format!("Failed to emit result: {e}"))?;

    Ok(())
}

fn join_segment_text(segments: &[Segment]) -> String {
    segments
        .iter()
        .map(|segment| segment.text.trim())
        .filter(|segment| !segment.is_empty())
        .collect::<Vec<_>>()
        .join(" ")
}

/// Transcribe an audio file by path. Decodes, resamples to 16kHz, runs Whisper.
#[tauri::command]
pub async fn transcribe_file(
    state: tauri::State<'_, AppState>,
    path: String,
    engine: Option<String>,
    model_id: Option<String>,
    language: String,
    initial_prompt: String,
    remove_fillers: bool,
    british_english: bool,
    anti_hallucination: bool,
    format_mode: String,
    profile_id: Option<String>,
) -> Result<serde_json::Value, String> {
    let resolved_profile_id =
        profile_id.unwrap_or_else(|| kon_storage::DEFAULT_PROFILE_ID.to_string());

    let profile = kon_storage::database::get_profile(&state.db, &resolved_profile_id)
        .await
        .map_err(|e| e.to_string())?
        .ok_or_else(|| format!("Profile {resolved_profile_id} not found"))?;

    let profile_terms: Vec<String> =
        kon_storage::database::list_profile_terms(&state.db, &resolved_profile_id)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .map(|t| t.term)
            .collect();

    let engine_name = engine.unwrap_or_else(|| "whisper".to_string());
    let model_id =
        model_id.unwrap_or_else(|| default_model_id_for_engine(&engine_name).to_string());
    // None: transcribe paths don't enforce sequential-GPU mode. That's
    // owned by the Settings-level load flows (see models.rs).
    ensure_model_loaded(&state, &engine_name, &model_id, None).await?;

    let engine = pick_engine(&state, &engine_name)?;
    let options = TranscriptionOptions {
        language: Some(language),
        initial_prompt: build_initial_prompt(
            &initial_prompt,
            &profile.initial_prompt,
            &profile_terms,
        ),
    };
    let engine_name_for_worker = engine_name.clone();

    let timed = tokio::task::spawn_blocking(move || {
        let audio = kon_audio::decode_audio_file(Path::new(&path)).map_err(|e| e.to_string())?;
        let resampled = kon_audio::resample_to_16khz(&audio).map_err(|e| e.to_string())?;
        transcribe_samples_sync(
            engine,
            &engine_name_for_worker,
            resampled.into_samples(),
            options,
        )
    })
    .await
    .map_err(|e| e.to_string())??;

    let dictionary_terms = profile_terms.clone();

    let mut segments: Vec<Segment> = timed.transcript.segments().to_vec();
    let raw_text = join_segment_text(&segments);
    post_process_segments(
        &mut segments,
        &PostProcessOptions {
            remove_fillers,
            british_english,
            anti_hallucination,
            format_mode: FormatMode::parse(&format_mode),
            dictionary_terms,
        },
        Some(state.llm_engine.as_ref()),
    );

    Ok(serde_json::json!({
        "engine": engine_name,
        "modelId": model_id,
        "segments": segments,
        "language": timed.transcript.language(),
        "duration": timed.transcript.duration(),
        "inference_ms": timed.inference_ms,
        "raw_text": raw_text,
    }))
}

/// Transcribe raw PCM f32 samples (Parakeet). Emits "transcription-result" event.
#[tauri::command]
pub async fn transcribe_pcm_parakeet(
    state: tauri::State<'_, AppState>,
    app: tauri::AppHandle,
    samples: Vec<f32>,
    chunk_id: u32,
    remove_fillers: bool,
    british_english: bool,
    anti_hallucination: bool,
    format_mode: String,
    profile_id: Option<String>,
) -> Result<(), String> {
    let resolved_profile_id =
        profile_id.unwrap_or_else(|| kon_storage::DEFAULT_PROFILE_ID.to_string());

    // Validate the profile exists so parakeet and whisper behave identically
    // when a bogus id slips through from the frontend.
    kon_storage::database::get_profile(&state.db, &resolved_profile_id)
        .await
        .map_err(|e| e.to_string())?
        .ok_or_else(|| format!("Profile {resolved_profile_id} not found"))?;

    let profile_terms: Vec<String> =
        kon_storage::database::list_profile_terms(&state.db, &resolved_profile_id)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .map(|t| t.term)
            .collect();

    let engine = state.parakeet_engine.clone();
    let options = TranscriptionOptions::default();

    let timed = tokio::task::spawn_blocking(move || {
        transcribe_samples_sync(engine, "parakeet", samples, options)
    })
    .await
    .map_err(|e| e.to_string())??;

    let dictionary_terms = profile_terms.clone();

    let mut segments: Vec<Segment> = timed.transcript.segments().to_vec();
    let raw_text = join_segment_text(&segments);
    post_process_segments(
        &mut segments,
        &PostProcessOptions {
            remove_fillers,
            british_english,
            anti_hallucination,
            format_mode: FormatMode::parse(&format_mode),
            dictionary_terms,
        },
        Some(state.llm_engine.as_ref()),
    );

    app.emit(
        "transcription-result",
        serde_json::json!({
            "status": "transcription",
            "segments": segments,
            "language": timed.transcript.language(),
            "duration": timed.transcript.duration(),
            "chunk_id": chunk_id,
            "inference_ms": timed.inference_ms,
            "raw_text": raw_text,
        }),
    )
    .map_err(|e| format!("Failed to emit result: {e}"))?;

    Ok(())
}