searchi/apps/desktop/src-tauri/src/semantic_commands.rs

//! Desktop-owned orchestration: managed model setup and persistent hybrid queries.
use crate::{model_download, AppState, Filters};
use searchi_core::SearchResult;
use searchi_service::{EmbeddingReport, HybridSearch};
use serde::Serialize;
use std::{
    path::PathBuf,
    sync::{
        atomic::{AtomicBool, AtomicU64, Ordering},
        Mutex,
    },
};
use tauri::State;

pub struct SemanticState {
    pub model_directory: PathBuf,
    pub worker: PathBuf,
    verified: AtomicBool,
    downloading: AtomicBool,
    embedding: AtomicBool,
    querying: AtomicBool,
    warmed: AtomicBool,
    warming: AtomicBool,
    downloaded_bytes: AtomicU64,
    error: Mutex<Option<String>>,
    service: Mutex<Option<HybridSearch>>,
}

impl SemanticState {
    pub fn new(model_directory: PathBuf, worker: PathBuf) -> Self {
        Self {
            model_directory,
            worker,
            verified: AtomicBool::new(false),
            downloading: AtomicBool::new(false),
            embedding: AtomicBool::new(false),
            querying: AtomicBool::new(false),
            warmed: AtomicBool::new(false),
            warming: AtomicBool::new(false),
            downloaded_bytes: AtomicU64::new(0),
            error: Mutex::new(None),
            service: Mutex::new(None),
        }
    }

    pub fn reset_session(&self) -> Result<(), String> {
        *self
            .service
            .try_lock()
            .map_err(|_| "Wait for the current semantic query before changing resources.")? = None;
        self.warmed.store(false, Ordering::Release);
        Ok(())
    }

    fn prepare_query_session(&self, database: PathBuf, threads: usize) -> Result<(), String> {
        self.warming.store(true, Ordering::Release);
        let _guard = WorkGuard(&self.warming);
        if !self.worker.is_file() {
            return Err(
                "The installed embedding worker is missing. Repair or reinstall Searchi.".into(),
            );
        }
        let mut service =
            HybridSearch::open(database, &self.model_directory, &self.worker, threads)
                .map_err(|error| error.to_string())?;
        service.warmup().map_err(|error| error.to_string())?;
        *self
            .service
            .lock()
            .map_err(|_| "Semantic session is unavailable.")? = Some(service);
        self.warmed.store(true, Ordering::Release);
        Ok(())
    }

    /// First-launch setup is authorized product behavior; inference remains offline.
    pub fn begin_automatic_setup(self: &std::sync::Arc<Self>, database: PathBuf, threads: usize) {
        if self
            .downloading
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_err()
        {
            return;
        }
        let semantic = self.clone();
        tauri::async_runtime::spawn_blocking(move || {
            let _guard = WorkGuard(&semantic.downloading);
            semantic.record_error(None);
            let result =
                model_download::provision(&semantic.model_directory, &semantic.downloaded_bytes);
            semantic.verified.store(result.is_ok(), Ordering::Release);
            let result = result.and_then(|_| semantic.prepare_query_session(database, threads));
            semantic.record_error(result.err());
        });
    }

    fn record_error(&self, value: Option<String>) {
        if let Ok(mut error) = self.error.lock() {
            *error = value;
        }
    }
}

#[derive(Serialize)]
pub struct SemanticStatus {
    model: &'static str,
    ready: bool,
    model_directory: String,
    worker_ready: bool,
    downloading: bool,
    embedding: bool,
    device_policy: &'static str,
    error: Option<String>,
    downloaded_bytes: u64,
    total_bytes: u64,
    generation: Option<i64>,
    chunks: usize,
    stale: bool,
    pending: usize,
    warming: bool,
    worker_info: Option<searchi_service::WorkerInfo>,
}

struct WorkGuard<'a>(&'a AtomicBool);
impl Drop for WorkGuard<'_> {
    fn drop(&mut self) {
        self.0.store(false, Ordering::Release);
    }
}

#[tauri::command]
pub fn semantic_status(state: State<'_, AppState>) -> Result<SemanticStatus, String> {
    let semantic = &state.semantic;
    let worker_ready = semantic.worker.is_file();
    let index =
        HybridSearch::embedding_status(&state.database).map_err(|error| error.to_string())?;
    Ok(SemanticStatus {
        model: "Google EmbeddingGemma 2",
        ready: worker_ready && semantic.warmed.load(Ordering::Acquire),
        model_directory: semantic.model_directory.to_string_lossy().into_owned(),
        worker_ready,
        downloading: semantic.downloading.load(Ordering::Acquire),
        embedding: semantic.embedding.load(Ordering::Acquire),
        device_policy: if std::env::var("SEARCHI_FORCE_CPU").as_deref() == Ok("1") {
            "cpu"
        } else {
            "auto"
        },
        error: semantic.error.lock().ok().and_then(|value| value.clone()),
        downloaded_bytes: semantic.downloaded_bytes.load(Ordering::Acquire),
        total_bytes: model_download::TOTAL_BYTES,
        generation: index.generation,
        chunks: index.chunks,
        stale: index.stale,
        pending: index.pending,
        warming: semantic.warming.load(Ordering::Acquire),
        worker_info: semantic
            .service
            .try_lock()
            .ok()
            .and_then(|service| service.as_ref().and_then(HybridSearch::worker_info)),
    })
}

#[tauri::command]
pub async fn setup_model(state: State<'_, AppState>) -> Result<(), String> {
    let semantic = state.semantic.clone();
    let database = state.database.clone();
    let threads = state
        .settings
        .lock()
        .map_err(|_| "Settings are unavailable.")?
        .value
        .cpu_threads;
    if semantic
        .downloading
        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
        .is_err()
    {
        return Err("Google model setup is already running.".into());
    }
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = WorkGuard(&semantic.downloading);
        semantic.record_error(None);
        let result =
            model_download::provision(&semantic.model_directory, &semantic.downloaded_bytes);
        semantic.verified.store(result.is_ok(), Ordering::Release);
        let result = result.and_then(|_| semantic.prepare_query_session(database, threads));
        semantic.record_error(result.as_ref().err().cloned());
        result
    })
    .await
    .map_err(|error| error.to_string())?
}

#[tauri::command]
pub async fn index_embeddings(state: State<'_, AppState>) -> Result<EmbeddingReport, String> {
    let semantic = state.semantic.clone();
    if !semantic.verified.load(Ordering::Acquire) || !semantic.worker.is_file() {
        return Err("Set up the Google model and native worker first.".into());
    }
    let threads = state
        .settings
        .lock()
        .map_err(|_| "Settings are unavailable.")?
        .value
        .cpu_threads;
    if semantic
        .embedding
        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
        .is_err()
    {
        return Err("Semantic indexing is already running.".into());
    }
    let database = state.database.clone();
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = WorkGuard(&semantic.embedding);
        semantic.record_error(None);
        // Indexing uses its own isolated worker; interactive queries keep their own session.
        let result = HybridSearch::open(
            database,
            &semantic.model_directory,
            &semantic.worker,
            threads,
        )
        .and_then(|mut service| service.index_embeddings())
        .map_err(|error| error.to_string());
        semantic.record_error(result.as_ref().err().cloned());
        result
    })
    .await
    .map_err(|error| error.to_string())?
}

#[tauri::command]
pub async fn hybrid_search(
    query: String,
    filters: Option<Filters>,
    state: State<'_, AppState>,
) -> Result<Vec<SearchResult>, String> {
    let semantic = state.semantic.clone();
    if !semantic.warmed.load(Ordering::Acquire) {
        return Err(
            "Semantic search is preparing in the background; keyword results are available.".into(),
        );
    }
    let database = state.database.clone();
    let threads = state
        .settings
        .lock()
        .map_err(|_| "Settings are unavailable.")?
        .value
        .cpu_threads;
    if semantic
        .querying
        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
        .is_err()
    {
        return Err("Semantic worker is busy; keyword results remain available.".into());
    }
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = WorkGuard(&semantic.querying);
        let mut cached = semantic
            .service
            .lock()
            .map_err(|_| "Semantic session is unavailable.")?;
        if cached.is_none() {
            *cached = Some(
                HybridSearch::open(
                    database,
                    &semantic.model_directory,
                    &semantic.worker,
                    threads,
                )
                .map_err(|error| error.to_string())?,
            );
        }
        cached
            .as_mut()
            .ok_or("Semantic session is unavailable.")?
            .search(&query, 50, &filters.unwrap_or_default().into())
            .map_err(|error| error.to_string())
    })
    .await
    .map_err(|error| error.to_string())?
}