minni 0.1.1

use crate::cli::init::ensure_initialized;
use crate::db::Database;
use crate::models::downloader::{
    get_dense_model_dir, get_model_dir, is_dense_model_installed, is_model_installed,
    DENSE_MODEL_ID, MODEL_ID,
};
use crate::search::{
    read_metadata, write_metadata, AnnIndex, BM25SearchResult, DenseRetriever, HybridSearch,
    IndexMetadata,
};
use anyhow::{anyhow, Context, Result};
use colored::Colorize;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::env;
use std::path::Path;

// ---------------------------------------------------------------------------
// Query file schema
// ---------------------------------------------------------------------------

/// A single expected result within a benchmark query.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ExpectedResult {
    /// File path (relative to project root, substring-matched).
    pub file_path: String,
    /// Optional symbol name (substring-matched, case-insensitive).
    pub symbol_name: Option<String>,
}

/// One entry in the benchmark query file.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct BenchmarkQuery {
    /// Natural-language search string.
    pub query: String,
    /// Results that must appear in the top-k to count as a hit.
    pub expected: Vec<ExpectedResult>,
    /// Recall@k cutoff (default 10).
    #[serde(default = "default_k")]
    pub k: usize,
}

fn default_k() -> usize {
    10
}

// ---------------------------------------------------------------------------
// Per-query evaluation result
// ---------------------------------------------------------------------------

#[derive(Debug, Serialize)]
pub struct QueryResult {
    pub query: String,
    pub k: usize,
    /// Number of expected results that were found in top-k.
    pub hits: usize,
    /// Total expected results.
    pub expected: usize,
    /// Recall@k = hits / expected  (0.0 – 1.0)
    pub recall_at_k: f64,
    /// MRR contribution: 1/rank of first hit, or 0 if no hit.
    pub reciprocal_rank: f64,
    /// Rank of the first matching result (1-based), or None.
    pub first_hit_rank: Option<usize>,
}

// ---------------------------------------------------------------------------
// Machine-readable JSON output
// ---------------------------------------------------------------------------

#[derive(Debug, Serialize)]
pub struct EvalJsonOutput {
    pub queries: Vec<QueryResult>,
    pub aggregate: AggregateMetrics,
}

#[derive(Debug, Serialize)]
pub struct AggregateMetrics {
    pub query_count: usize,
    pub mean_recall_at_k: f64,
    pub mean_reciprocal_rank: f64,
}

// ---------------------------------------------------------------------------
// Entry point
// ---------------------------------------------------------------------------

pub fn run(file: &Path, json: bool, min_mrr: Option<f64>) -> Result<()> {
    let project_root = env::current_dir()?;

    // Auto-initialize if not done
    ensure_initialized(&project_root, true)?;

    let db = Database::open(&project_root)?;

    let stats = db.get_stats()?;
    if stats.chunk_count == 0 {
        println!(
            "{} No indexed content found. Run {} first.",
            "!".yellow(),
            "minni index".cyan()
        );
        return Ok(());
    }

    // Load benchmark queries
    let raw = std::fs::read_to_string(file)
        .with_context(|| format!("Cannot read benchmark file: {}", file.display()))?;
    let queries: Vec<BenchmarkQuery> =
        serde_json::from_str(&raw).context("Invalid benchmark file format")?;

    if queries.is_empty() {
        println!("{} Benchmark file contains no queries.", "!".yellow());
        return Ok(());
    }

    // Validate each benchmark entry
    for bq in &queries {
        if bq.k < 1 {
            return Err(anyhow!(
                "Invalid benchmark entry for query {:?}: k must be >= 1 (got {})",
                bq.query,
                bq.k
            ));
        }
        if bq.expected.is_empty() {
            return Err(anyhow!(
                "Invalid benchmark entry for query {:?}: expected list must not be empty",
                bq.query
            ));
        }
    }

    // Build search engine (BM25 only — no model download required for CI)
    let bm25_index_path = project_root.join(".minni").join("bm25_index");
    let meta_path = bm25_index_path.join("minni_meta.json");
    let ann_index_path = project_root.join(".minni").join("ann_index.json");

    let model_path = if is_model_installed() {
        get_model_dir().ok()
    } else {
        None
    };
    let dense_model_path = if is_dense_model_installed() {
        get_dense_model_dir().ok()
    } else {
        None
    };

    ensure_bm25_compatible(&bm25_index_path, &meta_path, &db)?;

    let embeddings = db.get_all_embeddings().unwrap_or_default();
    let chunk_lookup = build_chunk_lookup(&db)?;
    let ann_index = AnnIndex::load(&ann_index_path).unwrap_or(None);

    let mut search = HybridSearch::new(
        &bm25_index_path,
        model_path.as_deref(),
        dense_model_path.as_deref(),
        ann_index,
        embeddings,
        chunk_lookup,
    )
    .context("Failed to initialize search")?;

    // Run evaluation
    if !json {
        println!("{} Running eval on {} queries…", "→".blue(), queries.len());
        println!();
    }

    let mut results: Vec<QueryResult> = Vec::with_capacity(queries.len());

    for bq in &queries {
        let k = bq.k;
        let (search_results, _meta) = search
            .search(&bq.query, k)
            .with_context(|| format!("Search failed for query: {}", bq.query))?;

        let (hits, first_hit_rank) = score_query(&bq.expected, &search_results, k);
        let expected = bq.expected.len().max(1); // guard against empty expected list
        let recall_at_k = hits as f64 / expected as f64;
        let reciprocal_rank = first_hit_rank.map(|r| 1.0 / r as f64).unwrap_or(0.0);

        results.push(QueryResult {
            query: bq.query.clone(),
            k,
            hits,
            expected: bq.expected.len(),
            recall_at_k,
            reciprocal_rank,
            first_hit_rank,
        });
    }

    // Aggregate metrics
    let n = results.len() as f64;
    let mean_recall = results.iter().map(|r| r.recall_at_k).sum::<f64>() / n;
    let mrr = results.iter().map(|r| r.reciprocal_rank).sum::<f64>() / n;

    let aggregate = AggregateMetrics {
        query_count: results.len(),
        mean_recall_at_k: mean_recall,
        mean_reciprocal_rank: mrr,
    };

    if json {
        let output = EvalJsonOutput {
            queries: results,
            aggregate,
        };
        println!("{}", serde_json::to_string_pretty(&output)?);
    } else {
        print_table(&results, &aggregate);
    }

    // Non-zero exit if below threshold
    if let Some(threshold) = min_mrr {
        if mrr < threshold {
            eprintln!(
                "{} MRR {:.4} is below threshold {:.4}",
                "✗".red(),
                mrr,
                threshold
            );
            std::process::exit(1);
        }
    }

    Ok(())
}

// ---------------------------------------------------------------------------
// Scoring helpers
// ---------------------------------------------------------------------------

/// Check whether a search result matches an expected entry.
/// `file_path` is matched as a substring; `symbol_name` (when set) is
/// matched as a case-insensitive substring of the result's symbol name.
fn result_matches(result: &crate::search::SearchResult, expected: &ExpectedResult) -> bool {
    if !result.file_path.contains(&expected.file_path) {
        return false;
    }
    if let Some(expected_sym) = &expected.symbol_name {
        let expected_lower = expected_sym.to_lowercase();
        let result_sym = result.symbol_name.as_deref().unwrap_or("").to_lowercase();
        if !result_sym.contains(&expected_lower) {
            return false;
        }
    }
    true
}

/// Returns (hits, first_hit_rank).
/// `hits` is the number of distinct expected entries found in the top-k results.
/// `first_hit_rank` is the 1-based rank of the earliest matching result.
fn score_query(
    expected: &[ExpectedResult],
    results: &[crate::search::SearchResult],
    k: usize,
) -> (usize, Option<usize>) {
    let top_k = &results[..results.len().min(k)];
    let mut matched = vec![false; expected.len()];
    let mut first_hit_rank: Option<usize> = None;

    for (rank, result) in top_k.iter().enumerate() {
        let rank1 = rank + 1;
        for (j, exp) in expected.iter().enumerate() {
            if !matched[j] && result_matches(result, exp) {
                matched[j] = true;
                if first_hit_rank.is_none() {
                    first_hit_rank = Some(rank1);
                }
            }
        }
    }

    let hits = matched.iter().filter(|&&v| v).count();
    (hits, first_hit_rank)
}

// ---------------------------------------------------------------------------
// Human-readable table
// ---------------------------------------------------------------------------

fn print_table(results: &[QueryResult], aggregate: &AggregateMetrics) {
    // Column widths
    let query_w = results
        .iter()
        .map(|r| r.query.len())
        .max()
        .unwrap_or(5)
        .max(5)
        .min(50);

    // Header
    println!(
        "{:<qw$}  {:>4}  {:>9}  {:>9}  {:>9}",
        "Query",
        "k",
        "Recall@k",
        "RR",
        "1st Hit",
        qw = query_w,
    );
    println!("{}", "─".repeat(query_w + 40));

    for r in results {
        let truncated_query = if r.query.len() > query_w {
            format!("{}…", &r.query[..query_w - 1])
        } else {
            r.query.clone()
        };

        let recall_str = format!("{:.3}", r.recall_at_k);
        let rr_str = format!("{:.3}", r.reciprocal_rank);
        let hit_str = r
            .first_hit_rank
            .map(|h| h.to_string())
            .unwrap_or_else(|| "-".to_string());

        let recall_colored = if r.recall_at_k >= 0.8 {
            recall_str.green()
        } else if r.recall_at_k >= 0.4 {
            recall_str.yellow()
        } else {
            recall_str.red()
        };

        let rr_colored = if r.reciprocal_rank >= 0.5 {
            rr_str.green()
        } else if r.reciprocal_rank > 0.0 {
            rr_str.yellow()
        } else {
            rr_str.red()
        };

        println!(
            "{:<qw$}  {:>4}  {:>9}  {:>9}  {:>9}",
            truncated_query,
            r.k,
            recall_colored,
            rr_colored,
            hit_str,
            qw = query_w,
        );
    }

    println!("{}", "─".repeat(query_w + 40));

    // Aggregate row
    let mrr_str = format!("{:.3}", aggregate.mean_reciprocal_rank);
    let recall_str = format!("{:.3}", aggregate.mean_recall_at_k);

    let mrr_colored = if aggregate.mean_reciprocal_rank >= 0.5 {
        mrr_str.green()
    } else if aggregate.mean_reciprocal_rank > 0.0 {
        mrr_str.yellow()
    } else {
        mrr_str.red()
    };

    println!(
        "{:<qw$}  {:>4}  {:>9}  {:>9}",
        "AGGREGATE".bold(),
        format!("n={}", aggregate.query_count),
        recall_str.bold().to_string(),
        mrr_colored,
        qw = query_w,
    );
    println!();
    println!(
        "Mean Recall@k: {}   MRR: {}   Queries: {}",
        format!("{:.3}", aggregate.mean_recall_at_k).cyan(),
        format!("{:.3}", aggregate.mean_reciprocal_rank).cyan(),
        aggregate.query_count.to_string().cyan(),
    );
}

// ---------------------------------------------------------------------------
// Search-init helpers (mirrors search.rs)
// ---------------------------------------------------------------------------

fn ensure_bm25_compatible(bm25_index_path: &Path, meta_path: &Path, db: &Database) -> Result<()> {
    use crate::search::BM25Index;

    let model_signature = format!("{}+{}", MODEL_ID, DENSE_MODEL_ID);
    let expected = IndexMetadata::expected(BM25Index::schema_signature()?, &model_signature);
    let current = read_metadata(meta_path)?;
    let needs_rebuild = match current {
        Some(meta) => !meta.matches(&expected),
        None => true,
    };

    if !needs_rebuild {
        return Ok(());
    }

    println!(
        "{} Rebuilding BM25 index due to schema/model changes…",
        "→".blue()
    );
    if bm25_index_path.exists() {
        std::fs::remove_dir_all(bm25_index_path)?;
    }
    let bm25_index = BM25Index::new(bm25_index_path)?;
    let chunks = db.get_all_chunks()?;
    bm25_index.index_chunks(&chunks)?;
    if is_dense_model_installed() {
        if let Ok(dense_model_path) = get_dense_model_dir() {
            let mut dense = DenseRetriever::new(&dense_model_path)?;
            let texts: Vec<String> = chunks
                .iter()
                .map(|c| format!("{}\n{}", c.file_path, c.content))
                .collect();
            let vectors = dense.embed_texts(&texts)?;
            db.clear_embeddings()?;
            for (chunk, vector) in chunks.iter().zip(vectors.iter()) {
                db.upsert_embedding(&chunk.id, vector)?;
            }
        }
    }
    write_metadata(meta_path, &expected)?;
    Ok(())
}

fn build_chunk_lookup(db: &Database) -> Result<HashMap<String, BM25SearchResult>> {
    let chunks = db.get_all_chunks()?;
    let mut lookup = HashMap::with_capacity(chunks.len());
    for chunk in chunks {
        lookup.insert(
            chunk.id.clone(),
            BM25SearchResult {
                chunk_id: chunk.id,
                file_path: chunk.file_path,
                content: chunk.content,
                start_line: chunk.start_line as usize,
                end_line: chunk.end_line as usize,
                chunk_type: chunk.chunk_type,
                language: chunk.language,
                symbol_name: chunk.symbol_name,
                score: 0.0,
                parent_symbol: chunk.parent_symbol,
                signature: chunk.signature,
                doc_comment: chunk.doc_comment,
                module_path: chunk.module_path,
            },
        );
    }
    Ok(lookup)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn make_result(file_path: &str, symbol_name: Option<&str>) -> SearchResult {
        SearchResult {
            chunk_id: "test".to_string(),
            file_path: file_path.to_string(),
            content: String::new(),
            start_line: 1,
            end_line: 10,
            chunk_type: "function".to_string(),
            language: "rust".to_string(),
            symbol_name: symbol_name.map(|s| s.to_string()),
            score: 1.0,
            parent_symbol: None,
            signature: None,
            doc_comment: None,
            module_path: None,
        }
    }

    #[test]
    fn test_result_matches_file_only() {
        let result = make_result("src/search/bm25.rs", Some("search"));
        let expected = ExpectedResult {
            file_path: "search/bm25.rs".to_string(),
            symbol_name: None,
        };
        assert!(result_matches(&result, &expected));
    }

    #[test]
    fn test_result_matches_file_and_symbol() {
        let result = make_result("src/search/bm25.rs", Some("BM25Index"));
        let expected = ExpectedResult {
            file_path: "bm25.rs".to_string(),
            symbol_name: Some("bm25index".to_string()),
        };
        assert!(result_matches(&result, &expected));
    }

    #[test]
    fn test_result_no_match_wrong_file() {
        let result = make_result("src/search/dense.rs", Some("search"));
        let expected = ExpectedResult {
            file_path: "bm25.rs".to_string(),
            symbol_name: None,
        };
        assert!(!result_matches(&result, &expected));
    }

    #[test]
    fn test_result_no_match_wrong_symbol() {
        let result = make_result("src/search/bm25.rs", Some("other_fn"));
        let expected = ExpectedResult {
            file_path: "bm25.rs".to_string(),
            symbol_name: Some("search".to_string()),
        };
        assert!(!result_matches(&result, &expected));
    }

    #[test]
    fn test_score_query_first_hit() {
        let expected = vec![ExpectedResult {
            file_path: "search/bm25.rs".to_string(),
            symbol_name: None,
        }];
        let results = vec![
            make_result("src/other.rs", None),
            make_result("src/search/bm25.rs", None),
        ];
        let (hits, first_rank) = score_query(&expected, &results, 10);
        assert_eq!(hits, 1);
        assert_eq!(first_rank, Some(2));
    }

    #[test]
    fn test_score_query_miss() {
        let expected = vec![ExpectedResult {
            file_path: "search/bm25.rs".to_string(),
            symbol_name: None,
        }];
        let results = vec![make_result("src/other.rs", None)];
        let (hits, first_rank) = score_query(&expected, &results, 10);
        assert_eq!(hits, 0);
        assert_eq!(first_rank, None);
    }

    #[test]
    fn test_score_query_partial_hits() {
        let expected = vec![
            ExpectedResult {
                file_path: "search/bm25.rs".to_string(),
                symbol_name: None,
            },
            ExpectedResult {
                file_path: "search/dense.rs".to_string(),
                symbol_name: None,
            },
        ];
        let results = vec![
            make_result("src/search/bm25.rs", None),
            make_result("src/other.rs", None),
        ];
        let (hits, first_rank) = score_query(&expected, &results, 10);
        assert_eq!(hits, 1);
        assert_eq!(first_rank, Some(1));
    }

    #[test]
    fn test_score_query_respects_k() {
        let expected = vec![ExpectedResult {
            file_path: "search/bm25.rs".to_string(),
            symbol_name: None,
        }];
        let results = vec![
            make_result("src/other.rs", None),
            make_result("src/other2.rs", None),
            make_result("src/search/bm25.rs", None), // rank 3, outside k=2
        ];
        let (hits, first_rank) = score_query(&expected, &results, 2);
        assert_eq!(hits, 0);
        assert_eq!(first_rank, None);
    }

    #[test]
    fn test_default_k() {
        assert_eq!(default_k(), 10);
    }

    #[test]
    fn test_benchmark_query_deserialization() {
        let json = r#"{"query":"test","expected":[{"file_path":"foo.rs"}]}"#;
        let bq: BenchmarkQuery = serde_json::from_str(json).unwrap();
        assert_eq!(bq.k, 10);
        assert_eq!(bq.expected[0].symbol_name, None);
    }
}