aurora_semantic/engine/

mod.rs

//! Engine module providing the main public API.
//!
//! The `Engine` struct is the primary entry point for using aurora-semantic.

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

use dashmap::DashMap;

use crate::chunker::{Chunker, DefaultChunker};
use crate::config::{EngineConfig, WorkspaceConfig};
use crate::embeddings::{create_embedding_text, Embedder, HashEmbedder};
use crate::error::{Error, Result};
use crate::ignore::FileWalker;
use crate::search::{LexicalIndex, SearchQuery, SemanticIndex};
use crate::storage::{DiskStorage, Storage, WorkspaceMetadata};
use crate::types::{
    Chunk, Document, DocumentId, IndexPhase, IndexProgress, Language, SearchResult, WorkspaceId,
    WorkspaceStats,
};

/// Progress callback type.
pub type ProgressCallback = Box<dyn Fn(IndexProgress) + Send + Sync>;

/// The main semantic search engine.
pub struct Engine {
    config: EngineConfig,
    storage: Arc<DiskStorage>,
    chunker: Arc<DefaultChunker>,
    embedder: Arc<dyn Embedder>,
    workspaces: DashMap<WorkspaceId, WorkspaceState>,
}

/// State of an indexed workspace.
struct WorkspaceState {
    #[allow(dead_code)]
    metadata: WorkspaceMetadata,
    documents: Vec<Document>,
    chunks: Vec<Chunk>,
    lexical_index: Arc<LexicalIndex>,
    semantic_index: Arc<SemanticIndex>,
}

impl Engine {
    /// Create a new engine with hash-based embeddings (for testing).
    pub fn new(config: EngineConfig) -> Result<Self> {
        let storage = Arc::new(DiskStorage::new(config.index_dir.clone())?);
        let chunker = Arc::new(DefaultChunker::new(config.chunking.clone())?);
        let embedder: Arc<dyn Embedder> = Arc::new(HashEmbedder::new(config.embedding.dimension));

        Ok(Self {
            config,
            storage,
            chunker,
            embedder,
            workspaces: DashMap::new(),
        })
    }

    /// Create an engine with a custom embedder (e.g., ONNX model).
    pub fn with_embedder<E: Embedder + 'static>(config: EngineConfig, embedder: E) -> Result<Self> {
        let storage = Arc::new(DiskStorage::new(config.index_dir.clone())?);
        let chunker = Arc::new(DefaultChunker::new(config.chunking.clone())?);

        Ok(Self {
            config,
            storage,
            chunker,
            embedder: Arc::new(embedder),
            workspaces: DashMap::new(),
        })
    }

    /// Get the engine configuration.
    pub fn config(&self) -> &EngineConfig {
        &self.config
    }

    /// Index a workspace with an auto-generated ID.
    pub async fn index_workspace(
        &self,
        workspace_config: WorkspaceConfig,
        progress: Option<ProgressCallback>,
    ) -> Result<WorkspaceId> {
        self.index_workspace_with_id(workspace_config, None, progress).await
    }

    /// Index a workspace with an optional custom ID.
    /// If `workspace_id` is None, a new UUID will be generated.
    /// If `workspace_id` is Some, the provided ID will be used (useful for re-indexing).
    pub async fn index_workspace_with_id(
        &self,
        workspace_config: WorkspaceConfig,
        workspace_id: Option<WorkspaceId>,
        progress: Option<ProgressCallback>,
    ) -> Result<WorkspaceId> {
        let workspace_id = workspace_id.unwrap_or_else(WorkspaceId::new);
        let root_path = workspace_config.root_path.clone();

        if !root_path.exists() {
            return Err(Error::DirectoryNotFound(root_path));
        }

        // Initialize storage
        self.storage.init_workspace(&workspace_id, &root_path)?;

        // Report scanning phase
        if let Some(ref cb) = progress {
            cb(IndexProgress {
                phase: IndexPhase::Scanning,
                processed: 0,
                total: 0,
                current_file: None,
                eta_seconds: None,
            });
        }

        // Walk and collect files
        let walker = FileWalker::new(root_path.clone(), self.config.ignore.clone())?;
        let files = walker.walk()?;
        let total_files = files.len();

        // Parse files into documents and chunks
        if let Some(ref cb) = progress {
            cb(IndexProgress {
                phase: IndexPhase::Parsing,
                processed: 0,
                total: total_files,
                current_file: None,
                eta_seconds: None,
            });
        }

        let mut documents = Vec::new();
        let mut chunks = Vec::new();

        for (i, file) in files.iter().enumerate() {
            if let Some(ref cb) = progress {
                cb(IndexProgress {
                    phase: IndexPhase::Parsing,
                    processed: i,
                    total: total_files,
                    current_file: Some(file.clone()),
                    eta_seconds: None,
                });
            }

            if let Ok((doc, file_chunks)) = self.process_file(file, &root_path) {
                documents.push(doc);
                chunks.extend(file_chunks);
            }
        }

        // Generate embeddings
        if let Some(ref cb) = progress {
            cb(IndexProgress {
                phase: IndexPhase::Embedding,
                processed: 0,
                total: chunks.len(),
                current_file: None,
                eta_seconds: None,
            });
        }

        let mut embeddings = Vec::new();
        let batch_size = self.config.embedding.batch_size;

        for (batch_idx, batch) in chunks.chunks(batch_size).enumerate() {
            if let Some(ref cb) = progress {
                cb(IndexProgress {
                    phase: IndexPhase::Embedding,
                    processed: batch_idx * batch_size,
                    total: chunks.len(),
                    current_file: None,
                    eta_seconds: None,
                });
            }

            let texts: Vec<String> = batch.iter().map(create_embedding_text).collect();
            let text_refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();

            if let Ok(batch_embeddings) = self.embedder.embed_batch(&text_refs) {
                for (chunk, embedding) in batch.iter().zip(batch_embeddings) {
                    embeddings.push((chunk.id.0.to_string(), embedding));
                }
            }
        }

        // Build indexes
        if let Some(ref cb) = progress {
            cb(IndexProgress {
                phase: IndexPhase::Indexing,
                processed: 0,
                total: 1,
                current_file: None,
                eta_seconds: None,
            });
        }

        // Build lexical index
        let lexical_path = self.storage.tantivy_index_path(&workspace_id);
        let lexical_index = LexicalIndex::create(&lexical_path, self.config.search.clone())?;

        let doc_map: HashMap<String, &Document> =
            documents.iter().map(|d| (d.id.0.to_string(), d)).collect();
        lexical_index.index_chunks(&chunks, &doc_map)?;

        // Build semantic index
        let mut semantic_index =
            SemanticIndex::new(self.embedder.dimension(), self.config.search.clone())?;
        semantic_index.add_embeddings(&embeddings)?;

        // Save to storage
        if let Some(ref cb) = progress {
            cb(IndexProgress {
                phase: IndexPhase::Persisting,
                processed: 0,
                total: 1,
                current_file: None,
                eta_seconds: None,
            });
        }

        self.storage.save_documents(&workspace_id, &documents)?;
        self.storage.save_chunks(&workspace_id, &chunks)?;
        self.storage.save_embeddings(&workspace_id, &embeddings)?;

        let vector_path = self.storage.vector_index_path(&workspace_id);
        semantic_index.save(&vector_path)?;

        // Update metadata
        let mut metadata = self.storage.load_workspace_metadata(&workspace_id)?;
        for doc in &documents {
            let chunk_count = chunks.iter().filter(|c| c.document_id == doc.id).count();
            metadata.record_document(
                doc.relative_path.clone(),
                doc.id.clone(),
                doc.content_hash.clone(),
                doc.size_bytes,
                doc.language,
                chunk_count,
            );
        }
        self.storage
            .save_workspace_metadata(&workspace_id, &metadata)?;

        // Store in memory
        self.workspaces.insert(
            workspace_id.clone(),
            WorkspaceState {
                metadata,
                documents,
                chunks,
                lexical_index: Arc::new(lexical_index),
                semantic_index: Arc::new(semantic_index),
            },
        );

        if let Some(ref cb) = progress {
            cb(IndexProgress {
                phase: IndexPhase::Complete,
                processed: total_files,
                total: total_files,
                current_file: None,
                eta_seconds: Some(0.0),
            });
        }

        Ok(workspace_id)
    }

    /// Load an existing workspace index.
    pub fn load_workspace(&self, workspace_id: &WorkspaceId) -> Result<()> {
        if !self.storage.workspace_exists(workspace_id) {
            return Err(Error::WorkspaceNotFound(workspace_id.clone()));
        }

        let metadata = self.storage.load_workspace_metadata(workspace_id)?;
        let documents = self.storage.load_documents(workspace_id)?;
        let chunks = self.storage.load_chunks(workspace_id)?;

        let lexical_path = self.storage.tantivy_index_path(workspace_id);
        let lexical_index = LexicalIndex::open(&lexical_path, self.config.search.clone())?;

        let vector_path = self.storage.vector_index_path(workspace_id);
        let semantic_index = SemanticIndex::load(
            &vector_path,
            self.embedder.dimension(),
            self.config.search.clone(),
        )?;

        self.workspaces.insert(
            workspace_id.clone(),
            WorkspaceState {
                metadata,
                documents,
                chunks,
                lexical_index: Arc::new(lexical_index),
                semantic_index: Arc::new(semantic_index),
            },
        );

        Ok(())
    }

    /// Search for code across a workspace.
    pub fn search(
        &self,
        workspace_id: &WorkspaceId,
        query: SearchQuery,
    ) -> Result<Vec<SearchResult>> {
        let state = self
            .workspaces
            .get(workspace_id)
            .ok_or_else(|| Error::WorkspaceNotFound(workspace_id.clone()))?;

        // Build lookup maps
        let chunk_map: HashMap<String, &Chunk> = state
            .chunks
            .iter()
            .map(|c| (c.id.0.to_string(), c))
            .collect();

        let doc_map: HashMap<String, &Document> = state
            .documents
            .iter()
            .map(|d| (d.id.0.to_string(), d))
            .collect();

        match query.mode {
            crate::config::SearchMode::Lexical => {
                state.lexical_index.search_index(&query, &chunk_map, &doc_map)
            }
            crate::config::SearchMode::Semantic => {
                state.semantic_index.search_with_embedder(
                    &query,
                    self.embedder.as_ref(),
                    &chunk_map,
                    &doc_map,
                )
            }
            crate::config::SearchMode::Hybrid => {
                let lexical_results =
                    state.lexical_index.search_index(&query, &chunk_map, &doc_map)?;
                let semantic_results = state.semantic_index.search_with_embedder(
                    &query,
                    self.embedder.as_ref(),
                    &chunk_map,
                    &doc_map,
                )?;

                // Merge results
                let weights = vec![
                    self.config.search.lexical_weight,
                    self.config.search.semantic_weight,
                ];
                Ok(crate::search::merge_results(
                    vec![lexical_results, semantic_results],
                    &weights,
                    query.limit,
                ))
            }
        }
    }

    /// Search with a simple text query.
    pub fn search_text(
        &self,
        workspace_id: &WorkspaceId,
        text: &str,
    ) -> Result<Vec<SearchResult>> {
        let query = SearchQuery::new(text)
            .limit(self.config.search.default_limit)
            .min_score(self.config.search.min_score);

        self.search(workspace_id, query)
    }

    /// Delete a workspace and all its data.
    pub fn delete_workspace(&self, workspace_id: &WorkspaceId) -> Result<()> {
        self.workspaces.remove(workspace_id);
        self.storage.delete_workspace(workspace_id)
    }

    /// List all indexed workspaces.
    pub fn list_workspaces(&self) -> Result<Vec<WorkspaceId>> {
        self.storage.list_workspaces()
    }

    /// Get statistics for a workspace.
    pub fn get_workspace_stats(&self, workspace_id: &WorkspaceId) -> Result<WorkspaceStats> {
        self.storage.get_workspace_stats(workspace_id)
    }

    /// Find a workspace by its root path.
    /// Returns the workspace ID if found, None otherwise.
    pub fn find_workspace_by_path(&self, root_path: &Path) -> Result<Option<WorkspaceId>> {
        let workspaces = self.list_workspaces()?;
        let normalized_path = root_path.to_string_lossy().to_lowercase().replace('\\', "/");
        
        for ws_id in workspaces {
            if let Ok(stats) = self.get_workspace_stats(&ws_id) {
                let indexed_path = stats.root_path.to_string_lossy().to_lowercase().replace('\\', "/");
                if indexed_path == normalized_path {
                    return Ok(Some(ws_id));
                }
            }
        }
        Ok(None)
    }

    /// Index or re-index a workspace by path.
    /// If the workspace was previously indexed, it will be deleted and re-indexed with the same ID.
    pub async fn index_or_reindex_workspace(
        &self,
        workspace_config: WorkspaceConfig,
        progress: Option<ProgressCallback>,
    ) -> Result<WorkspaceId> {
        let root_path = workspace_config.root_path.clone();
        
        // Check if workspace already exists
        let existing_id = self.find_workspace_by_path(&root_path)?;
        
        if let Some(ws_id) = existing_id {
            // Delete existing workspace first
            self.delete_workspace(&ws_id)?;
            // Re-index with the same ID
            self.index_workspace_with_id(workspace_config, Some(ws_id), progress).await
        } else {
            // New workspace, generate new ID
            self.index_workspace(workspace_config, progress).await
        }
    }

    fn process_file(&self, path: &Path, root: &Path) -> Result<(Document, Vec<Chunk>)> {
        let content = std::fs::read_to_string(path)?;
        let metadata = std::fs::metadata(path)?;
        let hash = blake3::hash(content.as_bytes()).to_hex().to_string();

        let extension = path.extension().and_then(|e| e.to_str()).unwrap_or("");
        let language = Language::from_extension(extension);
        let relative_path = path.strip_prefix(root).unwrap_or(path).to_path_buf();

        let document = Document {
            id: DocumentId::new(),
            relative_path,
            absolute_path: path.to_path_buf(),
            language,
            content_hash: hash,
            size_bytes: metadata.len(),
            modified_at: metadata
                .modified()
                .ok()
                .map(chrono::DateTime::from)
                .unwrap_or_else(chrono::Utc::now),
        };

        let chunks = self.chunker.chunk(&document, &content)?;
        Ok((document, chunks))
    }
}

unsafe impl Send for Engine {}
unsafe impl Sync for Engine {}