chronara/src-tauri/src/lib.rs

//! Chronara - Meeting transcription and summarization using local AI models.
//!
//! This is a pure Rust backend using:
//! - whisper-rs for transcription
//! - llama-cpp-2 for summarization
//! - voice_activity_detector for speaker separation

use parking_lot::Mutex;
use std::path::PathBuf;
use std::sync::Arc;
use tauri::{Emitter, Manager, State};
use tracing::{debug, info};

pub mod ml;
pub mod storage;

use ml::summarizer::{get_model_filename, LlamaSummarizer};
use ml::transcriber::{TranscriptSegment, WhisperTranscriber};
use ml::vad::SpeakerSeparator;
use ml::audio::AudioCapture;
use storage::{RecordingStorage, StoredRecording};

/// Application state containing the ML models and audio capture.
struct AppState {
    transcriber: Mutex<WhisperTranscriber>,
    summarizer: Mutex<Option<LlamaSummarizer>>,
    speaker_separator: Mutex<Option<SpeakerSeparator>>,
    audio_capture: Mutex<Option<AudioCapture>>,
    logs: Arc<Mutex<Vec<String>>>,
    storage: Mutex<Option<RecordingStorage>>,
}

impl AppState {
    fn new() -> Self {
        Self {
            transcriber: Mutex::new(WhisperTranscriber::new()),
            summarizer: Mutex::new(None),
            speaker_separator: Mutex::new(None),
            audio_capture: Mutex::new(None),
            logs: Arc::new(Mutex::new(Vec::new())),
            storage: Mutex::new(None),
        }
    }
}

/// Emit a log message to the frontend.
fn emit_log(app_handle: &tauri::AppHandle, logs: &Arc<Mutex<Vec<String>>>, message: &str) {
    {
        let mut logs_guard = logs.lock();
        logs_guard.push(message.to_string());
        if logs_guard.len() > 100 {
            logs_guard.remove(0);
        }
    }
    let _ = app_handle.emit("backend-log", message);
    info!("{}", message);
}

/// Get the models directory based on environment.
fn get_models_dir(app_handle: &tauri::AppHandle) -> PathBuf {
    // Production mode - use app data directory (user-writable)
    // On Windows: %APPDATA%\com.chronara.app\models
    // On macOS: ~/Library/Application Support/com.chronara.app/models
    // On Linux: ~/.local/share/com.chronara.app/models
    if let Ok(app_data_dir) = app_handle.path().app_data_dir() {
        return app_data_dir.join("models");
    }

    // Fallback: Development mode - use project models directory
    let current_dir = std::env::current_dir().unwrap_or_default();
    let project_root = if current_dir.ends_with("src-tauri") {
        current_dir.parent().unwrap().to_path_buf()
    } else {
        current_dir
    };
    project_root.join("models")
}

/// Check if the required models exist.
#[tauri::command]
fn check_models(app_handle: tauri::AppHandle) -> Result<bool, String> {
    let models_dir = get_models_dir(&app_handle);

    // Check for LLaMA model
    let llama_model = models_dir.join(get_model_filename("3B"));
    let llama_exists = llama_model.exists();

    // Check for Whisper model
    let whisper_model = models_dir.join("whisper").join("ggml-base.bin");
    let whisper_exists = whisper_model.exists();

    debug!(
        "Models check: llama={} ({}), whisper={} ({})",
        llama_exists,
        llama_model.display(),
        whisper_exists,
        whisper_model.display()
    );

    // Both models are required
    Ok(llama_exists && whisper_exists)
}

/// Download a file from a URL with progress tracking.
async fn download_file(
    app_handle: &tauri::AppHandle,
    logs: &Arc<Mutex<Vec<String>>>,
    url: &str,
    dest_path: &std::path::Path,
    model_name: &str,
) -> Result<(), String> {
    use futures_util::StreamExt;
    use std::io::Write;

    emit_log(app_handle, logs, &format!("[Models] Downloading {}...", model_name));

    let client = reqwest::Client::builder()
        .redirect(reqwest::redirect::Policy::limited(10))
        .timeout(std::time::Duration::from_secs(3600))
        .connect_timeout(std::time::Duration::from_secs(30))
        .build()
        .map_err(|e| format!("Failed to create HTTP client: {}", e))?;

    let response = client
        .get(url)
        .send()
        .await
        .map_err(|e| format!("Failed to start download: {}", e))?;

    if !response.status().is_success() {
        return Err(format!("Download failed with status: {}", response.status()));
    }

    let total_size = response.content_length().unwrap_or(0);
    let total_mb = total_size as f64 / 1_048_576.0;

    emit_log(app_handle, logs, &format!("[Models] {} size: {:.1} MB", model_name, total_mb));

    // Ensure parent directory exists
    if let Some(parent) = dest_path.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| format!("Failed to create directory: {}", e))?;
    }

    // Download to a temp file first
    let temp_path = dest_path.with_extension("downloading");
    let mut file = std::fs::File::create(&temp_path)
        .map_err(|e| format!("Failed to create temp file: {}", e))?;

    let mut downloaded: u64 = 0;
    let mut last_progress_percent: u64 = 0;
    let mut stream = response.bytes_stream();

    while let Some(chunk_result) = stream.next().await {
        let chunk = chunk_result.map_err(|e| format!("Download error: {}", e))?;
        file.write_all(&chunk)
            .map_err(|e| format!("Failed to write to file: {}", e))?;

        downloaded += chunk.len() as u64;

        if total_size > 0 {
            let progress_percent = (downloaded * 100) / total_size;
            if progress_percent >= last_progress_percent + 10 {
                last_progress_percent = progress_percent;
                let downloaded_mb = downloaded as f64 / 1_048_576.0;
                emit_log(
                    app_handle,
                    logs,
                    &format!("[Models] {}: {:.1} MB / {:.1} MB ({}%)", model_name, downloaded_mb, total_mb, progress_percent),
                );
            }
        }
    }

    file.flush().map_err(|e| format!("Failed to flush file: {}", e))?;
    drop(file);

    std::fs::rename(&temp_path, dest_path)
        .map_err(|e| format!("Failed to finalize download: {}", e))?;

    emit_log(app_handle, logs, &format!("[Models] {} download complete!", model_name));
    Ok(())
}

/// Download the required models.
#[tauri::command]
async fn download_models(
    state: State<'_, AppState>,
    app_handle: tauri::AppHandle,
) -> Result<String, String> {
    let logs = Arc::clone(&state.logs);
    let models_dir = get_models_dir(&app_handle);

    std::fs::create_dir_all(&models_dir)
        .map_err(|e| format!("Failed to create models directory: {}", e))?;

    // Download LLaMA model if needed
    let llama_model = models_dir.join(get_model_filename("3B"));
    if !llama_model.exists() {
        let llama_url = "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q4_K_M.gguf";
        download_file(&app_handle, &logs, llama_url, &llama_model, "LLaMA 3.2 3B (~2GB)").await?;
    } else {
        emit_log(&app_handle, &logs, "[Models] LLaMA model already present");
    }

    // Download Whisper model if needed
    let whisper_dir = models_dir.join("whisper");
    let whisper_model = whisper_dir.join("ggml-base.bin");
    if !whisper_model.exists() {
        let whisper_url = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.bin";
        download_file(&app_handle, &logs, whisper_url, &whisper_model, "Whisper Base (~142MB)").await?;
    } else {
        emit_log(&app_handle, &logs, "[Models] Whisper model already present");
    }

    emit_log(&app_handle, &logs, "[Models] All models ready!");
    Ok("Models downloaded successfully".to_string())
}

/// Initialize the ML models.
#[tauri::command]
async fn initialize_models(
    state: State<'_, AppState>,
    app_handle: tauri::AppHandle,
) -> Result<String, String> {
    let logs = Arc::clone(&state.logs);
    let models_dir = get_models_dir(&app_handle);

    emit_log(&app_handle, &logs, "[Init] Initializing ML models...");

    // Initialize LLaMA summarizer
    let llama_model_path = models_dir.join(get_model_filename("3B"));
    if llama_model_path.exists() {
        emit_log(&app_handle, &logs, "[Init] Loading LLaMA model...");

        match LlamaSummarizer::new() {
            Ok(mut summarizer) => {
                if let Err(e) = summarizer.load_model(&llama_model_path) {
                    emit_log(&app_handle, &logs, &format!("[Init ERROR] Failed to load LLaMA: {}", e));
                } else {
                    *state.summarizer.lock() = Some(summarizer);
                    emit_log(&app_handle, &logs, "[Init] LLaMA model loaded successfully");
                }
            }
            Err(e) => {
                emit_log(&app_handle, &logs, &format!("[Init ERROR] Failed to create summarizer: {}", e));
            }
        }
    } else {
        emit_log(&app_handle, &logs, "[Init WARNING] LLaMA model not found, summarization disabled");
    }

    // Initialize Whisper transcriber (lazy load on first use)
    let whisper_model_path = models_dir.join("whisper").join("ggml-base.bin");
    if whisper_model_path.exists() {
        emit_log(&app_handle, &logs, "[Init] Loading Whisper model...");

        let mut transcriber = state.transcriber.lock();
        if let Err(e) = transcriber.load_model(&whisper_model_path) {
            emit_log(&app_handle, &logs, &format!("[Init ERROR] Failed to load Whisper: {}", e));
        } else {
            emit_log(&app_handle, &logs, "[Init] Whisper model loaded successfully");
        }
    } else {
        emit_log(&app_handle, &logs, "[Init] Whisper model not found, will download on first transcription");
    }

    // Initialize VAD for speaker separation
    emit_log(&app_handle, &logs, "[Init] Initializing voice activity detector...");
    match SpeakerSeparator::new() {
        Ok(separator) => {
            *state.speaker_separator.lock() = Some(separator);
            emit_log(&app_handle, &logs, "[Init] VAD initialized successfully");
        }
        Err(e) => {
            emit_log(&app_handle, &logs, &format!("[Init WARNING] VAD initialization failed: {}", e));
        }
    }

    // Initialize storage
    emit_log(&app_handle, &logs, "[Init] Initializing recording storage...");
    if let Ok(app_data_dir) = app_handle.path().app_data_dir() {
        match RecordingStorage::new(&app_data_dir) {
            Ok(storage) => {
                *state.storage.lock() = Some(storage);
                emit_log(&app_handle, &logs, "[Init] Recording storage initialized successfully");
            }
            Err(e) => {
                emit_log(&app_handle, &logs, &format!("[Init WARNING] Storage initialization failed: {}", e));
            }
        }
    }

    emit_log(&app_handle, &logs, "[Init] Model initialization complete");
    Ok("Models initialized".to_string())
}

/// Start recording audio.
#[tauri::command]
fn start_recording(
    state: State<'_, AppState>,
    app_handle: tauri::AppHandle,
) -> Result<String, String> {
    let logs = Arc::clone(&state.logs);

    emit_log(&app_handle, &logs, "[Audio] Starting recording...");

    let mut audio_guard = state.audio_capture.lock();

    // Create audio capture if not exists
    if audio_guard.is_none() {
        match AudioCapture::new() {
            Ok(capture) => {
                *audio_guard = Some(capture);
            }
            Err(e) => {
                let msg = format!("[Audio ERROR] Failed to create audio capture: {}", e);
                emit_log(&app_handle, &logs, &msg);
                return Err(msg);
            }
        }
    }

    // Start recording
    if let Some(ref mut capture) = *audio_guard {
        if let Err(e) = capture.start_recording() {
            let msg = format!("[Audio ERROR] Failed to start recording: {}", e);
            emit_log(&app_handle, &logs, &msg);
            return Err(msg);
        }
    }

    emit_log(&app_handle, &logs, "[Audio] Recording started");
    Ok("Recording started".to_string())
}

/// Stop recording (for real-time mode, transcription already happened during recording).
#[tauri::command]
async fn stop_recording(
    state: State<'_, AppState>,
    app_handle: tauri::AppHandle,
) -> Result<String, String> {
    let logs = Arc::clone(&state.logs);

    emit_log(&app_handle, &logs, "[Audio] Stopping recording...");

    // Stop the recording and get the duration
    let duration = {
        let mut audio_guard = state.audio_capture.lock();
        if let Some(ref mut capture) = *audio_guard {
            let samples = capture.stop_recording();
            samples.len() as f32 / 16000.0
        } else {
            return Err("No active recording".to_string());
        }
    };

    emit_log(&app_handle, &logs, &format!("[Audio] Recording stopped. Total duration: {:.1}s", duration));
    Ok("Recording stopped".to_string())
}

/// Stop recording and transcribe all at once (batch mode).
#[tauri::command]
async fn stop_recording_batch(
    state: State<'_, AppState>,
    app_handle: tauri::AppHandle,
) -> Result<Vec<TranscriptSegment>, String> {
    let logs = Arc::clone(&state.logs);

    emit_log(&app_handle, &logs, "[Audio] Stopping recording (batch mode)...");

    // Get the audio samples
    let audio_samples = {
        let mut audio_guard = state.audio_capture.lock();
        if let Some(ref mut capture) = *audio_guard {
            capture.stop_recording()
        } else {
            return Err("No active recording".to_string());
        }
    };

    let duration = audio_samples.len() as f32 / 16000.0;
    emit_log(&app_handle, &logs, &format!("[Audio] Captured {:.1}s of audio", duration));

    if audio_samples.is_empty() {
        return Err("No audio captured".to_string());
    }

    // Transcribe the audio
    emit_log(&app_handle, &logs, "[Transcribe] Starting transcription...");

    let app_handle_clone = app_handle.clone();
    let mut segments = {
        let transcriber = state.transcriber.lock();
        if !transcriber.is_loaded() {
            emit_log(&app_handle, &logs, "[Transcribe ERROR] Whisper model not loaded");
            return Err("Whisper model not loaded. Please ensure the model is downloaded.".to_string());
        }

        transcriber.transcribe_with_progress(&audio_samples, move |progress| {
            // Emit progress event to frontend
            let _ = app_handle_clone.emit("transcription-progress", progress);
        })
            .map_err(|e| format!("Transcription failed: {}", e))?
    };

    emit_log(&app_handle, &logs, &format!("[Transcribe] Got {} segments", segments.len()));

    // Apply speaker labels using VAD
    if let Some(ref mut separator) = *state.speaker_separator.lock() {
        emit_log(&app_handle, &logs, "[Speaker] Applying speaker labels...");
        segments = separator.apply_speaker_labels(&audio_samples, segments)
            .map_err(|e| format!("Speaker separation failed: {}", e))?;
    }

    Ok(segments)
}

/// Transcribe a chunk of audio (for real-time transcription).
#[tauri::command]
async fn transcribe_chunk(
    state: State<'_, AppState>,
    audio_data: Vec<f32>,
    app_handle: tauri::AppHandle,
) -> Result<Vec<TranscriptSegment>, String> {
    let transcriber = state.transcriber.lock();

    if !transcriber.is_loaded() {
        return Err("Whisper model not loaded".to_string());
    }

    // Clone the app handle for the closure
    let app_handle_clone = app_handle.clone();

    let segments = transcriber.transcribe_with_progress(&audio_data, move |progress| {
        // Emit progress event to frontend
        let _ = app_handle_clone.emit("transcription-progress", progress);
    })
        .map_err(|e| format!("Transcription failed: {}", e))?;

    Ok(segments)
}

/// Get the next chunk of audio for real-time transcription.
/// Returns the audio chunk and the new offset to use for the next call.
#[tauri::command]
async fn get_audio_chunk(
    state: State<'_, AppState>,
    last_offset: usize,
) -> Result<(Vec<f32>, usize), String> {
    let audio_guard = state.audio_capture.lock();
    if let Some(ref capture) = *audio_guard {
        Ok(capture.extract_chunk(last_offset))
    } else {
        Err("No active recording".to_string())
    }
}

/// Get remaining audio without modifying the buffer (for final processing).
#[tauri::command]
async fn get_remaining_audio(
    state: State<'_, AppState>,
    last_offset: usize,
) -> Result<Vec<f32>, String> {
    let audio_guard = state.audio_capture.lock();
    if let Some(ref capture) = *audio_guard {
        Ok(capture.get_remaining_audio(last_offset))
    } else {
        Err("No active recording".to_string())
    }
}

/// Generate a summary from a transcript.
#[tauri::command]
async fn summarize(
    state: State<'_, AppState>,
    app_handle: tauri::AppHandle,
    transcript: String,
) -> Result<String, String> {
    let logs = Arc::clone(&state.logs);

    emit_log(&app_handle, &logs, "[Summary] Generating summary...");

    let summarizer_guard = state.summarizer.lock();
    let summarizer = summarizer_guard.as_ref()
        .ok_or("Summarizer not initialized")?;

    if !summarizer.is_loaded() {
        return Err("LLaMA model not loaded".to_string());
    }

    let app_handle_clone = app_handle.clone();
    let summary = summarizer.summarize_with_progress(&transcript, move |progress| {
        // Emit progress event to frontend
        let _ = app_handle_clone.emit("summary-progress", progress);
    })
        .map_err(|e| format!("Summarization failed: {}", e))?;

    emit_log(&app_handle, &logs, &format!("[Summary] Generated {} character summary", summary.len()));

    Ok(summary)
}

/// Get backend logs.
#[tauri::command]
fn get_backend_logs(state: State<'_, AppState>) -> Result<Vec<String>, String> {
    Ok(state.logs.lock().clone())
}

/// Check if models are loaded and ready.
#[tauri::command]
fn check_ready(state: State<'_, AppState>) -> Result<bool, String> {
    let summarizer = state.summarizer.lock();

    // At minimum, we need the summarizer loaded
    // Whisper can be loaded on first use
    let ready = summarizer.as_ref().map_or(false, |s| s.is_loaded());

    Ok(ready)
}

/// Save a recording to persistent storage.
#[tauri::command]
fn save_recording(
    state: State<'_, AppState>,
    recording: StoredRecording,
) -> Result<(), String> {
    let storage_guard = state.storage.lock();
    let storage = storage_guard.as_ref()
        .ok_or("Storage not initialized")?;

    storage.save_recording(&recording)
        .map_err(|e| format!("Failed to save recording: {}", e))?;

    Ok(())
}

/// Load all recordings from persistent storage.
#[tauri::command]
fn load_recordings(
    state: State<'_, AppState>,
) -> Result<Vec<StoredRecording>, String> {
    let storage_guard = state.storage.lock();
    let storage = storage_guard.as_ref()
        .ok_or("Storage not initialized")?;

    storage.load_all_recordings()
        .map_err(|e| format!("Failed to load recordings: {}", e))
}

/// Delete a recording from persistent storage.
#[tauri::command]
fn delete_recording(
    state: State<'_, AppState>,
    recording_id: String,
) -> Result<(), String> {
    let storage_guard = state.storage.lock();
    let storage = storage_guard.as_ref()
        .ok_or("Storage not initialized")?;

    storage.delete_recording(&recording_id)
        .map_err(|e| format!("Failed to delete recording: {}", e))?;

    Ok(())
}

/// Update a recording (e.g., to add summary).
#[tauri::command]
fn update_recording(
    state: State<'_, AppState>,
    recording: StoredRecording,
) -> Result<(), String> {
    let storage_guard = state.storage.lock();
    let storage = storage_guard.as_ref()
        .ok_or("Storage not initialized")?;

    storage.update_recording(&recording)
        .map_err(|e| format!("Failed to update recording: {}", e))?;

    Ok(())
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    // Initialize tracing
    tracing_subscriber::fmt::init();

    info!("Starting Chronara with native Rust backend");

    tauri::Builder::default()
        .plugin(tauri_plugin_opener::init())
        .manage(AppState::new())
        .invoke_handler(tauri::generate_handler![
            check_models,
            download_models,
            initialize_models,
            start_recording,
            stop_recording,
            transcribe_chunk,
            get_audio_chunk,
            get_remaining_audio,
            summarize,
            get_backend_logs,
            check_ready,
            save_recording,
            load_recordings,
            delete_recording,
            update_recording,
        ])
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}