//! Dense vector search coordinator
//!
//! Combines embeddings and HNSW index for semantic search.

use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};

use parking_lot::RwLock;

use crate::compaction::hnsw_rebuild::{HnswRebuilder, RebuildResult};
use crate::embeddings::{CandleEmbedder, Embedder, EmbeddingConfig};
use crate::error::{MemdError, Result};
use crate::index::{HnswConfig, HnswIndex};
use crate::metrics::IndexStats;
use crate::types::{ChunkId, TenantId};

/// Result of a dense search with chunk content
#[derive(Debug, Clone)]
pub struct DenseSearchResult {
    /// Chunk ID
    pub chunk_id: ChunkId,
    /// Cosine similarity score (0.0 to 1.0)
    pub score: f32,
}

/// Configuration for dense search
#[derive(Debug, Clone)]
pub struct DenseSearchConfig {
    /// HNSW configuration
    pub hnsw: HnswConfig,
    /// Whether to persist index
    pub persist: bool,
    // Temporarily removed during Candle migration
    // /// Embedding model to use
    // pub model: EmbeddingModel,
}

impl Default for DenseSearchConfig {
    fn default() -> Self {
        Self {
            hnsw: HnswConfig::default(),
            persist: true,
            // model: EmbeddingModel::default(),
        }
    }
}

fn embedding_batch_size() -> usize {
    std::env::var("MEMD_EMBED_BATCH_SIZE")
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .filter(|size| *size > 0)
        .unwrap_or(EmbeddingConfig::default().batch_size)
}

/// Dense search coordinator for a tenant
pub struct DenseSearcher {
    /// Embedding model (shared across tenants)
    embedder: Arc<dyn Embedder>,
    /// Per-tenant HNSW indices
    indices: RwLock<HashMap<String, Arc<HnswIndex>>>,
    /// Base path for index persistence
    base_path: Option<PathBuf>,
    /// Disable disk writes and index mutations.
    read_only: bool,
    /// Configuration
    config: DenseSearchConfig,
}

impl DenseSearcher {
    /// Create a new dense searcher with Candle embedder
    pub fn new(config: DenseSearchConfig) -> Result<Self> {
        let mut embedding_config = EmbeddingConfig::default();
        embedding_config.batch_size = embedding_batch_size();
        let embedder = Arc::new(CandleEmbedder::with_config(embedding_config)?);

        // Update HNSW config dimension to match model
        let mut updated_config = config.clone();
        updated_config.hnsw.dimension = embedder.dimension();

        Ok(Self {
            embedder,
            indices: RwLock::new(HashMap::new()),
            base_path: None,
            read_only: false,
            config: updated_config,
        })
    }

    /// Create with custom embedder (for testing with MockEmbedder)
    pub fn with_embedder(embedder: Arc<dyn Embedder>, config: DenseSearchConfig) -> Self {
        Self {
            embedder,
            indices: RwLock::new(HashMap::new()),
            base_path: None,
            read_only: false,
            config,
        }
    }

    /// Set base path for index persistence
    pub fn with_base_path(mut self, path: PathBuf) -> Self {
        self.base_path = Some(path);
        self
    }

    /// Configure read-only mode for persisted warm indexes.
    pub fn with_read_only(mut self, read_only: bool) -> Self {
        self.read_only = read_only;
        self
    }

    /// Get or create index for a tenant
    fn get_or_create_index(&self, tenant_id: &TenantId) -> Result<Arc<HnswIndex>> {
        let tenant_str = tenant_id.to_string();

        // Fast path: read lock
        {
            let indices = self.indices.read();
            if let Some(index) = indices.get(&tenant_str) {
                return Ok(Arc::clone(index));
            }
        }

        // Slow path: write lock + create
        let mut indices = self.indices.write();

        // Double-check
        if let Some(index) = indices.get(&tenant_str) {
            return Ok(Arc::clone(index));
        }

        let index = if self.config.persist {
            if let Some(ref base_path) = self.base_path {
                let index_path = base_path
                    .join("tenants")
                    .join(&tenant_str)
                    .join("warm_index");
                if self.read_only {
                    HnswIndex::with_persistence_read_only(self.config.hnsw.clone(), index_path)?
                } else {
                    HnswIndex::with_persistence(self.config.hnsw.clone(), index_path)?
                }
            } else {
                HnswIndex::new(self.config.hnsw.clone())
            }
        } else {
            HnswIndex::new(self.config.hnsw.clone())
        };

        let index = Arc::new(index);
        indices.insert(tenant_str, Arc::clone(&index));

        Ok(index)
    }

    /// Index a chunk embedding
    pub async fn index_chunk(
        &self,
        tenant_id: &TenantId,
        chunk_id: &ChunkId,
        text: &str,
    ) -> Result<()> {
        if self.read_only {
            return Err(MemdError::ReadOnlyStore {
                op: "dense_index_chunk".to_string(),
            });
        }
        let embedding = self.embedder.embed_query(text).await?;
        let index = self.get_or_create_index(tenant_id)?;
        index.insert(chunk_id, &embedding)?;

        tracing::debug!(
            tenant_id = %tenant_id,
            chunk_id = %chunk_id,
            "indexed chunk in HNSW"
        );

        Ok(())
    }

    /// Index multiple chunks in batch
    pub async fn index_batch(
        &self,
        tenant_id: &TenantId,
        chunks: &[(ChunkId, String)],
    ) -> Result<()> {
        if self.read_only {
            return Err(MemdError::ReadOnlyStore {
                op: "dense_index_batch".to_string(),
            });
        }
        if chunks.is_empty() {
            return Ok(());
        }

        // Batch embed all texts
        let texts: Vec<&str> = chunks.iter().map(|(_, text)| text.as_str()).collect();
        let embeddings = self.embedder.embed_texts(&texts).await?;

        // Insert into index
        let index = self.get_or_create_index(tenant_id)?;
        let items: Vec<(ChunkId, Vec<f32>)> = chunks
            .iter()
            .zip(embeddings.into_iter())
            .map(|((chunk_id, _), emb)| (chunk_id.clone(), emb))
            .collect();

        index.insert_batch(&items)?;

        tracing::debug!(
            tenant_id = %tenant_id,
            count = chunks.len(),
            "indexed batch in HNSW"
        );

        Ok(())
    }

    /// Search for similar chunks
    pub async fn search(
        &self,
        tenant_id: &TenantId,
        query: &str,
        k: usize,
    ) -> Result<Vec<DenseSearchResult>> {
        let (results, _, _) = self.search_with_timing(tenant_id, query, k).await?;
        Ok(results)
    }

    /// Search for similar chunks with timing information
    ///
    /// Returns (results, embed_time, search_time)
    pub async fn search_with_timing(
        &self,
        tenant_id: &TenantId,
        query: &str,
        k: usize,
    ) -> Result<(Vec<DenseSearchResult>, Duration, Duration)> {
        let embed_start = Instant::now();
        let query_embedding = self.embedder.embed_query(query).await?;
        let embed_time = embed_start.elapsed();

        let search_start = Instant::now();
        let index = self.get_or_create_index(tenant_id)?;
        let results = index.search(&query_embedding, k)?;
        let search_time = search_start.elapsed();

        let dense_results: Vec<DenseSearchResult> = results
            .into_iter()
            .map(|r| DenseSearchResult {
                chunk_id: r.chunk_id,
                score: r.score,
            })
            .collect();

        tracing::debug!(
            tenant_id = %tenant_id,
            query_len = query.len(),
            results = dense_results.len(),
            embed_ms = embed_time.as_millis(),
            search_ms = search_time.as_millis(),
            "dense search completed"
        );

        Ok((dense_results, embed_time, search_time))
    }

    /// Get index statistics for all tenants
    pub fn get_stats(&self) -> HashMap<String, IndexStats> {
        let indices = self.indices.read();
        let dimension = self.embedder.dimension();

        indices
            .iter()
            .map(|(tenant_id, index)| {
                let count = index.len() as u64;
                // Estimate memory: each embedding is dim * 4 bytes, plus HNSW overhead (~2x)
                let embedding_bytes = count * dimension as u64 * 4;
                let estimated_memory = embedding_bytes * 2;

                (
                    tenant_id.clone(),
                    IndexStats {
                        chunks_indexed: count,
                        embeddings_count: count,
                        embedding_dimension: dimension,
                        index_memory_bytes: estimated_memory,
                    },
                )
            })
            .collect()
    }

    /// Get index statistics for a specific tenant
    pub fn get_tenant_stats(&self, tenant_id: &TenantId) -> Option<IndexStats> {
        let indices = self.indices.read();
        let tenant_str = tenant_id.to_string();
        let dimension = self.embedder.dimension();

        indices.get(&tenant_str).map(|index| {
            let count = index.len() as u64;
            let embedding_bytes = count * dimension as u64 * 4;
            let estimated_memory = embedding_bytes * 2;

            IndexStats {
                chunks_indexed: count,
                embeddings_count: count,
                embedding_dimension: dimension,
                index_memory_bytes: estimated_memory,
            }
        })
    }

    /// Save all indices
    pub fn save_all(&self) -> Result<()> {
        if self.read_only {
            tracing::debug!("read-only dense searcher skips save_all");
            return Ok(());
        }
        let indices = self.indices.read();
        for (tenant_id, index) in indices.iter() {
            if let Err(e) = index.save() {
                tracing::warn!(tenant_id, error = %e, "failed to save index");
            }
        }
        Ok(())
    }

    /// Get embedding dimension
    pub fn dimension(&self) -> usize {
        self.embedder.dimension()
    }

    /// Search using a pre-computed query embedding (for tiered search)
    ///
    /// This avoids re-embedding when the caller already has the embedding.
    pub fn search_with_embedding(
        &self,
        tenant_id: &TenantId,
        query_embedding: &[f32],
        k: usize,
    ) -> Result<Vec<DenseSearchResult>> {
        let index = self.get_or_create_index(tenant_id)?;
        let results = index.search(query_embedding, k)?;

        Ok(results
            .into_iter()
            .map(|r| DenseSearchResult {
                chunk_id: r.chunk_id,
                score: r.score,
            })
            .collect())
    }

    /// Get the number of indexed chunks for a tenant
    pub fn index_len(&self, tenant_id: &TenantId) -> usize {
        let indices = self.indices.read();
        let tenant_str = tenant_id.to_string();
        indices.get(&tenant_str).map(|i| i.len()).unwrap_or(0)
    }

    /// Check whether a specific chunk is present in the tenant's HNSW
    /// mapping.
    ///
    /// Used by startup backfill: comparing counts (`index_len` vs metadata
    /// row count) is unreliable because HNSW's `next_id` counter never
    /// decrements on delete, so a tenant with deletes + re-adds + a crash
    /// can satisfy `index_len >= active_count` while still missing live
    /// chunks. Per-chunk membership is the authoritative signal.
    pub fn contains_chunk(&self, tenant_id: &TenantId, chunk_id: &ChunkId) -> bool {
        let indices = self.indices.read();
        let tenant_str = tenant_id.to_string();
        match indices.get(&tenant_str) {
            Some(index) => index
                .get_mapping()
                .read()
                .get_internal_id(chunk_id)
                .is_some(),
            None => false,
        }
    }

    /// Check whether any of the supplied chunks still have valid
    /// embeddings in the tenant cache.
    ///
    /// HNSW compaction keeps the chunk-to-internal-id mapping stable
    /// after rebuild, but invalidates excluded embeddings in the cache.
    /// This helper distinguishes "known mapping" from "still live in the
    /// dense rebuild source", which is the signal purge/compaction needs.
    pub(crate) fn has_valid_embeddings_for_chunks(
        &self,
        tenant_id: &TenantId,
        chunk_ids: &HashSet<ChunkId>,
    ) -> bool {
        if chunk_ids.is_empty() {
            return false;
        }

        let indices = self.indices.read();
        let tenant_str = tenant_id.to_string();
        let Some(index) = indices.get(&tenant_str) else {
            return false;
        };

        let mapping = index.get_mapping().read();
        let cache = index.get_embedding_cache().read();
        chunk_ids.iter().any(|chunk_id| {
            mapping
                .get_internal_id(chunk_id)
                .and_then(|internal_id| cache.get(internal_id))
                .is_some()
        })
    }

    /// Ensure the tenant's persisted HNSW index (mapping + embedding cache)
    /// is loaded into memory. The backfill calls this before deciding what is
    /// "missing", so a clean persisted index is not needlessly re-embedded.
    pub(crate) fn ensure_index_loaded(&self, tenant_id: &TenantId) -> Result<()> {
        self.get_or_create_index(tenant_id)?;
        Ok(())
    }

    /// Of `chunk_ids`, return those WITHOUT a live cached embedding in the
    /// loaded index — no mapping entry, OR a mapping entry whose embedding
    /// cache slot is empty (e.g. after a missing/corrupt `embeddings.bin`,
    /// where the loader keeps the mapping but resets the cache). Takes the
    /// per-tenant locks once. If the index is not loaded, every chunk is
    /// reported missing — callers should `ensure_index_loaded` first.
    ///
    /// This is the cache-aware membership the backfill needs: `contains_chunk`
    /// (mapping-only) would wrongly treat cache-less chunks as present and
    /// skip re-embedding them, leaving the index without usable vectors.
    pub(crate) fn chunks_missing_embeddings(
        &self,
        tenant_id: &TenantId,
        chunk_ids: &[ChunkId],
    ) -> Vec<ChunkId> {
        let indices = self.indices.read();
        let Some(index) = indices.get(&tenant_id.to_string()) else {
            return chunk_ids.to_vec();
        };
        let mapping = index.get_mapping().read();
        let cache = index.get_embedding_cache().read();
        chunk_ids
            .iter()
            .filter(|chunk_id| {
                mapping
                    .get_internal_id(chunk_id)
                    .and_then(|internal_id| cache.get(internal_id))
                    .is_none()
            })
            .cloned()
            .collect()
    }

    /// Embed a query text (exposes embedder for tiered search)
    pub async fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
        self.embedder.embed_query(text).await
    }

    /// Get rebuild stats (cache_size, index_size) for a tenant
    ///
    /// Returns (0, 0) if no index exists for the tenant.
    pub fn get_rebuild_stats(&self, tenant_id: &TenantId) -> (usize, usize) {
        let indices = self.indices.read();
        if let Some(index) = indices.get(&tenant_id.to_string()) {
            index.rebuild_stats()
        } else {
            (0, 0)
        }
    }

    /// Rebuild HNSW index for a tenant, excluding deleted chunk IDs
    ///
    /// Creates a new clean index from embedding cache, filtering out
    /// deleted entries. Returns rebuild statistics.
    ///
    /// Note: Full atomic swap requires HnswIndex to support construction
    /// from a pre-built Hnsw graph. Currently returns rebuild result
    /// only - the index is rebuilt in place conceptually.
    pub fn rebuild_hnsw_for_tenant(
        &self,
        tenant_id: &TenantId,
        deleted_chunk_ids: &HashSet<ChunkId>,
    ) -> Result<RebuildResult> {
        let tenant_str = tenant_id.to_string();

        // Get existing index (read lock)
        let indices = self.indices.read();
        let existing_index = indices.get(&tenant_str).ok_or_else(|| {
            MemdError::StorageError(format!("no index for tenant: {}", tenant_id))
        })?;

        // Convert chunk IDs to internal IDs
        let mapping = existing_index.get_mapping().read();
        let deleted_internal_ids: HashSet<usize> = deleted_chunk_ids
            .iter()
            .filter_map(|chunk_id| mapping.get_internal_id(chunk_id))
            .collect();
        drop(mapping);

        // Use HnswRebuilder to create a clean graph that excludes
        // the deleted internal IDs.
        let rebuilder = HnswRebuilder::new();
        let (new_hnsw, result) = rebuilder.rebuild_clean(
            existing_index,
            &deleted_internal_ids,
            existing_index.config(),
        )?;

        // Phase 3.2: swap the new graph into the live HnswIndex AND
        // mark the excluded IDs invalid in the embedding cache so
        // save/load stay consistent with the rebuilt graph. Without
        // the cache invalidation, a restart would resurrect the
        // compacted-away points (the Codex-flagged durability gap).
        existing_index.swap_graph(new_hnsw, &deleted_internal_ids);

        tracing::info!(
            tenant_id = %tenant_id,
            processed = result.embeddings_processed,
            included = result.embeddings_included,
            excluded = result.embeddings_excluded,
            duration_ms = result.duration.as_millis(),
            "HNSW rebuild completed and swapped for tenant"
        );

        Ok(result)
    }
}

impl Drop for DenseSearcher {
    fn drop(&mut self) {
        if self.config.persist {
            if let Err(e) = self.save_all() {
                tracing::warn!(error = %e, "failed to save indices on drop");
            }
        }
    }
}
