mirage-analyzer 1.9.1

//! Semantic search for Mirage — resolve natural language queries to function IDs
//! via HNSW embeddings (generated by Magellan, searched by Ollama).

use rusqlite::{Connection, OptionalExtension};
use serde::Deserialize;

/// Embedding provider configuration (read from `~/.config/magellan/config.toml`).
#[derive(Debug, Deserialize)]
struct MagellanConfig {
    embeddings: EmbeddingConfig,
}

#[derive(Debug, Deserialize)]
struct EmbeddingConfig {
    base_url: String,
    model: String,
    api_key: Option<String>,
}

/// Embed a natural-language query via the Ollama (or OpenAI-compatible) endpoint.
pub fn embed_query(
    base_url: &str,
    model: &str,
    api_key: &str,
    query: &str,
) -> anyhow::Result<Vec<f32>> {
    let url = format!("{}/api/embeddings", base_url.trim_end_matches('/'));

    #[derive(serde::Serialize)]
    struct OllamaEmbedRequest<'a> {
        model: &'a str,
        prompt: &'a str,
    }

    #[derive(serde::Deserialize)]
    struct OllamaEmbedResponse {
        embedding: Vec<f32>,
    }

    let body = OllamaEmbedRequest {
        model,
        prompt: query,
    };
    let json_body = serde_json::to_string(&body)
        .map_err(|e| anyhow::anyhow!("Failed to serialize embedding request: {}", e))?;

    let mut request = ureq::post(&url).header("Content-Type", "application/json");
    if !api_key.is_empty() {
        request = request.header("Authorization", &format!("Bearer {api_key}"));
    }
    let response = request
        .send(&json_body)
        .map_err(|e| anyhow::anyhow!("Ollama embedding request failed: {}", e))?;

    let body = response
        .into_body()
        .read_to_string()
        .map_err(|e| anyhow::anyhow!("Failed to read Ollama response body: {}", e))?;
    let resp: OllamaEmbedResponse = serde_json::from_str(&body)
        .map_err(|e| anyhow::anyhow!("Failed to parse Ollama embedding response: {}", e))?;

    Ok(resp.embedding)
}

/// Read `~/.config/magellan/config.toml` for embedding settings.
pub fn read_magellan_config() -> anyhow::Result<(String, String, String)> {
    let home = std::env::var("HOME")
        .or_else(|_| std::env::var("USERPROFILE"))
        .map_err(|_| anyhow::anyhow!("Could not determine home directory"))?;
    let config_path = std::path::Path::new(&home)
        .join(".config")
        .join("magellan")
        .join("config.toml");
    let contents = std::fs::read_to_string(&config_path)
        .map_err(|e| anyhow::anyhow!("Failed to read config at {:?}: {}", config_path, e))?;
    let config: MagellanConfig =
        toml::from_str(&contents).map_err(|e| anyhow::anyhow!("Failed to parse config: {}", e))?;
    let api_key = config.embeddings.api_key.unwrap_or_default();
    Ok((config.embeddings.base_url, config.embeddings.model, api_key))
}

/// Resolve a natural-language query to the top matching function entity ID.
///
/// 1. Reads Ollama config from `~/.config/magellan/config.toml`
/// 2. Embeds the query
/// 3. Loads the HNSW index from the database
/// 4. Returns the top entity ID (or `None` if no HNSW index exists)
pub fn resolve_semantic_query(
    db_path: &std::path::Path,
    query: &str,
    top_k: usize,
) -> anyhow::Result<Option<Vec<(i64, f32)>>> {
    let conn =
        Connection::open(db_path).map_err(|e| anyhow::anyhow!("Failed to open DB: {}", e))?;

    // Fast-fail if no HNSW index exists
    let hnsw_exists: bool = conn
        .query_row(
            "SELECT 1 FROM sqlite_master WHERE type='table' AND name='hnsw_indexes'",
            [],
            |_| Ok(true),
        )
        .unwrap_or(false);
    if !hnsw_exists {
        return Ok(None);
    }

    let (base_url, model, api_key) = read_magellan_config()?;
    let query_vector = embed_query(&base_url, &model, &api_key, query)?;

    let index = sqlitegraph::hnsw::HnswIndex::load_with_vectors(&conn, "symbols")
        .map_err(|e| anyhow::anyhow!("Failed to load HNSW index: {}", e))?;
    let results = index
        .search(&query_vector, top_k)
        .map_err(|e| anyhow::anyhow!("HNSW search failed: {}", e))?;

    let mut matches = Vec::new();
    for (vector_id, distance) in results {
        let metadata_json: Option<String> = conn
            .query_row(
                "SELECT metadata FROM hnsw_vectors WHERE id = ?1",
                [vector_id],
                |row| row.get(0),
            )
            .optional()
            .unwrap_or(None);
        let metadata_json = match metadata_json {
            Some(m) => m,
            None => continue,
        };
        let metadata: serde_json::Value =
            serde_json::from_str(&metadata_json).unwrap_or(serde_json::json!({}));
        if let Some(entity_id) = metadata["entity_id"].as_i64() {
            matches.push((entity_id, distance));
        }
    }

    Ok(Some(matches))
}