//! Desktop-owned orchestration: managed model setup and persistent hybrid queries.
use crate::{measured_outcome, model_download, AppState, Filters, MeasuredFailure};
use searchi_core::metrics::{MetricsMode, SearchResponse, TraceCollector};
use searchi_core::SearchResult;
use searchi_service::{
    EmbeddingReport, HybridSearch, InferenceBudget, InferenceRole, SemanticCoverage,
};
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,
    text_verified: AtomicBool,
    media_verified: AtomicBool,
    media_error: Mutex<Option<String>>,
    downloading: AtomicBool,
    embedding: AtomicBool,
    querying: AtomicBool,
    warmed: AtomicBool,
    warming: AtomicBool,
    downloaded_bytes: AtomicU64,
    error: Mutex<Option<String>>,
    service: Mutex<Option<HybridSearch>>,
    prepared_graphs: searchi_service::PreparedGraphs,
    inference_budget: InferenceBudget,
    lifecycle: Mutex<()>,
    media_tools: Option<searchi_media::FfmpegTools>,
    media_limits: searchi_media::MediaLimits,
    media_setup_error: Option<String>,
}

impl SemanticState {
    pub fn new(
        model_directory: PathBuf,
        worker: PathBuf,
        media_tools: Result<searchi_media::FfmpegTools, String>,
    ) -> Self {
        let (media_tools, media_setup_error) = match media_tools {
            Ok(tools) => (Some(tools), None),
            Err(error) => (None, Some(error)),
        };
        Self {
            model_directory,
            worker,
            verified: AtomicBool::new(false),
            text_verified: AtomicBool::new(false),
            media_verified: AtomicBool::new(false),
            media_error: Mutex::new(None),
            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),
            prepared_graphs: searchi_service::PreparedGraphs::default(),
            inference_budget: InferenceBudget::new(2),
            lifecycle: Mutex::new(()),
            media_tools,
            media_limits: searchi_media::MediaLimits::default(),
            media_setup_error,
        }
    }

    /// Keep generated video frames below the application's excluded data tree.
    pub fn with_media_scratch(mut self, directory: PathBuf) -> Self {
        self.media_limits.scratch_directory = Some(directory);
        self
    }

    fn ensure_idle(&self) -> Result<(), String> {
        if self.embedding.load(Ordering::Acquire)
            || self.downloading.load(Ordering::Acquire)
            || self.warming.load(Ordering::Acquire)
            || self.querying.load(Ordering::Acquire)
        {
            return Err("Wait for semantic work before changing resource settings.".into());
        }
        Ok(())
    }

    #[cfg(test)]
    fn reset_session_locked(&self) -> Result<(), String> {
        self.ensure_idle()?;
        *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(),
            );
        }
        self.inference_budget.configure(threads);
        let threads = self
            .inference_budget
            .worker_threads(InferenceRole::Foreground);
        let admission = self
            .inference_budget
            .acquire(InferenceRole::Foreground, threads)
            .map_err(|error| error.to_string())?;
        let mut service =
            HybridSearch::open_query(database, &self.model_directory, &self.worker, threads)
                .map_err(|error| error.to_string())?;
        drop(admission);
        service = service.with_budget(&self.inference_budget, InferenceRole::Foreground, threads);
        service
            .share_prepared_graphs(&self.prepared_graphs)
            .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) {
        let Ok(admission) = self.lifecycle.lock() else {
            return;
        };
        if self.embedding.load(Ordering::Acquire) || self.querying.load(Ordering::Acquire) {
            return;
        }
        if self
            .downloading
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_err()
        {
            return;
        }
        drop(admission);
        self.spawn_setup(database, threads);
    }

    fn spawn_setup(self: &std::sync::Arc<Self>, database: PathBuf, threads: usize) {
        let semantic = self.clone();
        tauri::async_runtime::spawn_blocking(move || {
            let _guard = WorkGuard(&semantic.downloading);
            let _ = semantic.run_setup(database, threads);
        });
    }

    fn run_setup(
        self: &std::sync::Arc<Self>,
        database: PathBuf,
        threads: usize,
    ) -> Result<(), String> {
        self.record_error(None);
        self.verified.store(false, Ordering::Release);
        self.text_verified.store(false, Ordering::Release);
        self.media_verified.store(false, Ordering::Release);
        if let Ok(mut error) = self.media_error.lock() {
            *error = None;
        }
        let mut warm = None;
        let provision = model_download::provision_with_verified_assets(
            &self.model_directory,
            &self.downloaded_bytes,
            |asset| match asset {
                searchi_assets::VerifiedAsset::TextBackbone => {
                    self.text_verified.store(true, Ordering::Release);
                    if !self.warmed.load(Ordering::Acquire) {
                        // Reserve warmup before publishing readiness to indexing admission.
                        self.warming.store(true, Ordering::Release);
                        let semantic = self.clone();
                        let database = database.clone();
                        warm = Some(
                            std::thread::Builder::new()
                                .name("searchi-query-prepare".into())
                                .spawn(move || {
                                    let result = semantic.prepare_query_session(database, threads);
                                    // Report immediately: the independent projector download can
                                    // continue for minutes after a failed query warmup.
                                    if let Err(error) = &result {
                                        semantic.record_error(Some(error.clone()));
                                    }
                                    result
                                })
                                .map_err(|error| {
                                    let error =
                                        format!("Could not start query preparation: {error}");
                                    self.record_error(Some(error.clone()));
                                    self.warming.store(false, Ordering::Release);
                                    error
                                }),
                        );
                    }
                }
                searchi_assets::VerifiedAsset::MediaProjector => {
                    self.media_verified.store(true, Ordering::Release);
                }
            },
        );
        self.verified.store(provision.is_ok(), Ordering::Release);
        let prepared = warm
            .map(|handle| {
                handle.and_then(|handle| {
                    handle
                        .join()
                        .unwrap_or_else(|_| Err("Query preparation failed unexpectedly.".into()))
                })
            })
            .unwrap_or(Ok(()));
        self.warming.store(false, Ordering::Release);
        self.finish_setup(provision, prepared)
    }

    fn finish_setup(
        &self,
        provision: Result<(), String>,
        prepared: Result<(), String>,
    ) -> Result<(), String> {
        if let Err(error) = &prepared {
            self.record_error(Some(error.clone()));
        }
        if let Err(error) = &provision {
            if self.text_verified.load(Ordering::Acquire) {
                if let Ok(mut media_error) = self.media_error.lock() {
                    *media_error = Some(error.clone());
                }
            } else {
                self.record_error(Some(error.clone()));
            }
        }
        prepared.and(provision)
    }

    /// Linearize resource reset, persisted settings and setup admission. The
    /// lifecycle lock never spans worker startup, inference or a service wait.
    pub fn change_resources(
        self: &std::sync::Arc<Self>,
        database: PathBuf,
        threads: usize,
        save: impl FnOnce() -> Result<(), String>,
    ) -> Result<(), String> {
        let admission = self
            .lifecycle
            .lock()
            .map_err(|_| "Semantic admission is unavailable.")?;
        self.ensure_idle()?;
        let mut service = self
            .service
            .try_lock()
            .map_err(|_| "Wait for the current semantic query before changing resources.")?;
        // A failed settings write must keep the existing warm session usable.
        save()?;
        *service = None;
        self.warmed.store(false, Ordering::Release);
        self.downloading.store(true, Ordering::Release);
        drop(service);
        drop(admission);
        self.spawn_setup(database, threads);
        Ok(())
    }

    fn desired_coverage(&self) -> SemanticCoverage {
        if self.verified.load(Ordering::Acquire) && self.media_tools.is_some() {
            SemanticCoverage::FullMedia
        } else {
            SemanticCoverage::TextOnly
        }
    }

    fn can_index(&self, text_indexing_allowed: bool) -> bool {
        self.text_verified.load(Ordering::Acquire)
            && self.worker.is_file()
            && self.warmed.load(Ordering::Acquire)
            && !self.warming.load(Ordering::Acquire)
            && !self.embedding.load(Ordering::Acquire)
            && (self.desired_coverage() == SemanticCoverage::FullMedia || text_indexing_allowed)
    }

    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,
    text_model_ready: bool,
    media_model_ready: bool,
    media_model_error: Option<String>,
    indexing_ready: bool,
    coverage: Option<SemanticCoverage>,
    requested_coverage: SemanticCoverage,
    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>,
    media_tools_ready: bool,
    media_setup_error: Option<String>,
    media_failures: Vec<searchi_service::MediaFailure>,
}

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 desired = semantic.desired_coverage();
    let index = HybridSearch::embedding_status_for_coverage(&state.database, desired)
        .map_err(|error| error.to_string())?;
    let indexing_ready = semantic.can_index(index.text_indexing_allowed);
    Ok(SemanticStatus {
        model: "Google EmbeddingGemma 2",
        ready: worker_ready && semantic.warmed.load(Ordering::Acquire),
        text_model_ready: semantic.text_verified.load(Ordering::Acquire),
        media_model_ready: semantic.media_verified.load(Ordering::Acquire),
        media_model_error: semantic
            .media_error
            .lock()
            .ok()
            .and_then(|error| error.clone()),
        indexing_ready,
        coverage: index.coverage,
        requested_coverage: desired,
        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)),
        media_tools_ready: semantic.media_tools.is_some(),
        media_setup_error: semantic.media_setup_error.clone(),
        media_failures: index.failures,
    })
}

#[tauri::command]
pub async fn setup_model(state: State<'_, AppState>) -> Result<(), String> {
    let semantic = state.semantic.clone();
    if semantic.embedding.load(Ordering::Acquire) || semantic.querying.load(Ordering::Acquire) {
        return Err("Wait for semantic work before restarting model setup.".into());
    }
    let database = state.database.clone();
    let threads = state
        .settings
        .lock()
        .map_err(|_| "Settings are unavailable.")?
        .value
        .cpu_threads;
    let admission = semantic
        .lifecycle
        .lock()
        .map_err(|_| "Semantic admission is unavailable.")?;
    if semantic.embedding.load(Ordering::Acquire) || semantic.querying.load(Ordering::Acquire) {
        return Err("Wait for semantic work before restarting model setup.".into());
    }
    if semantic
        .downloading
        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
        .is_err()
    {
        return Err("Google model setup is already running.".into());
    }
    drop(admission);
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = WorkGuard(&semantic.downloading);
        semantic.run_setup(database, threads)
    })
    .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();
    let admission = semantic
        .lifecycle
        .lock()
        .map_err(|_| "Semantic admission is unavailable.")?;
    let coverage = semantic.desired_coverage();
    let index = HybridSearch::embedding_status_for_coverage(&state.database, coverage)
        .map_err(|error| error.to_string())?;
    if !semantic.can_index(index.text_indexing_allowed) {
        return Err(
            "Semantic indexing is waiting for verified models and query preparation.".into(),
        );
    }
    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();
    drop(admission);
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = WorkGuard(&semantic.embedding);
        semantic.record_error(None);
        // Both workers share one clamped budget. Admission is per encoder call,
        // so waiting queries stop the next background batch without a service lock.
        let threads = semantic.inference_budget.worker_threads(InferenceRole::Background);
        let result = (|| {
        let admission = semantic.inference_budget.acquire(InferenceRole::Background, threads)?;
        let service = HybridSearch::open(
            &database,
            &semantic.model_directory,
            &semantic.worker,
            threads,
        )?;
        drop(admission);
        let mut service = service.with_budget(&semantic.inference_budget, InferenceRole::Background, threads);
            service.share_prepared_graphs(&semantic.prepared_graphs)?;
            match coverage {
                SemanticCoverage::TextOnly => service.index_text_embeddings_preserving_media()
                    .and_then(|report| report.ok_or_else(|| anyhow::anyhow!("Full media index retained while media setup finishes."))),
                SemanticCoverage::FullMedia => {
                    let tools = semantic.media_tools.as_ref().ok_or_else(|| anyhow::anyhow!("The media decoder is missing from this installation. Repair or reinstall searchi."))?;
                    service.index_embeddings_with_media(tools, &semantic.media_limits)
                }
            }
        })()
        .map_err(|error| error.to_string());
        // The publisher shares its prepared immutable ANN before activation.
        // Keep the query model and its cached vectors warm throughout indexing.
        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();
    let admission = semantic
        .lifecycle
        .lock()
        .map_err(|_| "Semantic admission is unavailable.")?;
    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();
    if semantic
        .querying
        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
        .is_err()
    {
        return Err("Semantic worker is busy; keyword results remain available.".into());
    }
    let pending = semantic.inference_budget.pending_foreground();
    drop(admission);
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = WorkGuard(&semantic.querying);
        let _pending = pending;
        let mut cached = semantic
            .service
            .lock()
            .map_err(|_| "Semantic session is unavailable.")?;
        if cached.is_none() {
            let threads = semantic
                .inference_budget
                .worker_threads(InferenceRole::Foreground);
            let admission = semantic
                .inference_budget
                .acquire(InferenceRole::Foreground, threads)
                .map_err(|error| error.to_string())?;
            let session = HybridSearch::open_query(
                database,
                &semantic.model_directory,
                &semantic.worker,
                threads,
            )
            .map_err(|error| error.to_string())?;
            drop(admission);
            *cached = Some(session.with_budget(
                &semantic.inference_budget,
                InferenceRole::Foreground,
                threads,
            ));
        }
        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())?
}

#[tauri::command]
pub async fn hybrid_search_measured(
    query: String,
    filters: Option<Filters>,
    debug_performance: bool,
    state: State<'_, AppState>,
) -> Result<SearchResponse<SearchResult>, MeasuredFailure> {
    let started = std::time::Instant::now();
    let mut trace = TraceCollector::new(if debug_performance {
        MetricsMode::Debug
    } else {
        MetricsMode::Basic
    });
    let semantic = state.semantic.clone();
    let admission = semantic.lifecycle.lock().map_err(|_| {
        MeasuredFailure::new(
            "Semantic admission is unavailable.",
            TraceCollector::new(MetricsMode::Basic),
        )
    })?;
    if !semantic.warmed.load(Ordering::Acquire) {
        return Err(MeasuredFailure::new(
            "Semantic search is preparing in the background; keyword results are available.",
            trace,
        ));
    }
    if semantic
        .querying
        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
        .is_err()
    {
        return Err(MeasuredFailure::new(
            "Semantic worker is busy; keyword results remain available.",
            trace,
        ));
    }
    let database = state.database.clone();
    let queue = trace.begin("desktop.worker_queue", None);
    let pending = semantic.inference_budget.pending_foreground();
    drop(admission);
    tauri::async_runtime::spawn_blocking(move || {
        let _guard = WorkGuard(&semantic.querying);
        let _pending = pending;
        trace.end(queue, Default::default());
        let result = (|| {
            let mut cached = trace.timed("desktop.semantic_lock", None, |_, _| {
                semantic
                    .service
                    .lock()
                    .map_err(|_| "Semantic session is unavailable.")
            })?;
            if cached.is_none() {
                let threads = semantic
                    .inference_budget
                    .worker_threads(InferenceRole::Foreground);
                let admission = semantic
                    .inference_budget
                    .acquire(InferenceRole::Foreground, threads)
                    .map_err(|error| error.to_string())?;
                let session = trace.timed("desktop.session_initialize", None, |_, _| {
                    HybridSearch::open_query(
                        database,
                        &semantic.model_directory,
                        &semantic.worker,
                        threads,
                    )
                    .map_err(|error| error.to_string())
                })?;
                drop(admission);
                *cached = Some(session.with_budget(
                    &semantic.inference_budget,
                    InferenceRole::Foreground,
                    threads,
                ));
            }
            cached
                .as_mut()
                .ok_or("Semantic session is unavailable.")?
                .search_traced(
                    &query,
                    50,
                    &filters.unwrap_or_default().into(),
                    &mut trace,
                    None,
                )
                .map_err(|error| error.to_string())
        })();
        measured_outcome(result, trace)
    })
    .await
    .map_err(|error| MeasuredFailure::task(error.to_string(), started))?
}

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

    #[test]
    fn text_admission_requires_verified_warm_session_and_preserves_full_index() {
        let runtime = tempfile::NamedTempFile::new().unwrap();
        let state = SemanticState::new(
            PathBuf::new(),
            runtime.path().into(),
            Err("no decoder".into()),
        );
        state.text_verified.store(true, Ordering::Release);
        assert!(!state.can_index(true));
        state.warmed.store(true, Ordering::Release);
        state.downloading.store(true, Ordering::Release);
        assert!(
            state.can_index(true),
            "projector download must not block verified text"
        );
        assert!(
            !state.can_index(false),
            "existing full generation must survive repair"
        );
        state.warming.store(true, Ordering::Release);
        assert!(!state.can_index(true));
        state.warming.store(false, Ordering::Release);
        state.embedding.store(true, Ordering::Release);
        assert!(!state.can_index(true));
        state.embedding.store(false, Ordering::Release);
        state.text_verified.store(false, Ordering::Release);
        assert!(
            !state.can_index(true),
            "cached warm model cannot authorize unverified indexing"
        );
    }

    #[test]
    fn full_coverage_requires_pair_verification_and_decoder() {
        let state = SemanticState::new(PathBuf::new(), PathBuf::new(), Err("no decoder".into()));
        state.verified.store(true, Ordering::Release);
        state.media_verified.store(true, Ordering::Release);
        assert_eq!(state.desired_coverage(), SemanticCoverage::TextOnly);
    }

    #[test]
    fn projector_error_preserves_warm_text_and_is_separate_from_query_error() {
        let state = SemanticState::new(PathBuf::new(), PathBuf::new(), Err("no decoder".into()));
        state.text_verified.store(true, Ordering::Release);
        state.warmed.store(true, Ordering::Release);
        assert!(state
            .finish_setup(Err("projector download failed".into()), Ok(()))
            .is_err());
        assert!(state.warmed.load(Ordering::Acquire));
        assert!(state.error.lock().unwrap().is_none());
        assert_eq!(
            state.media_error.lock().unwrap().as_deref(),
            Some("projector download failed")
        );
        assert!(!state.verified.load(Ordering::Acquire));
    }

    #[test]
    fn unverified_backbone_and_warmup_errors_remain_query_errors() {
        let state = SemanticState::new(PathBuf::new(), PathBuf::new(), Err("no decoder".into()));
        state
            .finish_setup(Err("backbone invalid".into()), Ok(()))
            .unwrap_err();
        assert_eq!(
            state.error.lock().unwrap().as_deref(),
            Some("backbone invalid")
        );
        assert!(state.media_error.lock().unwrap().is_none());
        state.text_verified.store(true, Ordering::Release);
        state
            .finish_setup(Ok(()), Err("worker startup failed".into()))
            .unwrap_err();
        assert_eq!(
            state.error.lock().unwrap().as_deref(),
            Some("worker startup failed")
        );
        assert!(!state.warmed.load(Ordering::Acquire));
    }

    #[test]
    fn failed_resource_save_preserves_ready_session_and_admission() {
        let state = std::sync::Arc::new(SemanticState::new(
            PathBuf::new(),
            PathBuf::new(),
            Err("not installed".into()),
        ));
        state.warmed.store(true, Ordering::Release);
        state.verified.store(true, Ordering::Release);
        let result = state.change_resources(PathBuf::new(), 2, || Err("disk write failed".into()));
        assert_eq!(result.unwrap_err(), "disk write failed");
        assert!(state.warmed.load(Ordering::Acquire));
        assert!(state.verified.load(Ordering::Acquire));
        assert!(!state.downloading.load(Ordering::Acquire));
        assert!(!state.warming.load(Ordering::Acquire));
        assert!(state.service.try_lock().is_ok());
        assert!(state.lifecycle.try_lock().is_ok());
    }

    #[test]
    fn resources_cannot_reset_active_semantic_work() {
        let state = SemanticState::new(PathBuf::new(), PathBuf::new(), Err("not installed".into()));
        for busy in [
            &state.embedding,
            &state.downloading,
            &state.warming,
            &state.querying,
        ] {
            state.warmed.store(true, Ordering::Release);
            busy.store(true, Ordering::Release);
            assert!(state.reset_session_locked().is_err());
            assert!(state.warmed.load(Ordering::Acquire));
            busy.store(false, Ordering::Release);
        }
        state.reset_session_locked().unwrap();
        assert!(!state.warmed.load(Ordering::Acquire));
    }
}