use std::collections::{BTreeMap, HashSet};

use anyhow::{anyhow, Context, Result};
use rusqlite::{params, Connection};

use super::types::{
    CategoryEvaluation, EvaluationLayers, EvidenceRef, GoldenDataset, GoldenEvalReport,
    GoldenMemory, GoldenQuery, MetricSums, QueryEvaluation, QueryMetrics, QueryStatus,
};

const RANK_K: usize = 10;

#[derive(Default)]
struct CategoryAccumulator {
    total_queries: usize,
    abstention_queries: usize,
    abstention_passed: usize,
    metrics: MetricSums,
}

pub fn load_dataset(dataset_path: &str) -> Result<GoldenDataset> {
    let content = std::fs::read_to_string(dataset_path)
        .with_context(|| format!("read golden eval dataset {dataset_path}"))?;
    let dataset: GoldenDataset = serde_json::from_str(&content)
        .with_context(|| format!("parse golden eval dataset {dataset_path}"))?;
    validate_dataset(&dataset)?;
    Ok(dataset)
}

pub fn run_dataset_path(
    conn: &Connection,
    dataset_path: &str,
    k: usize,
) -> Result<GoldenEvalReport> {
    let dataset = load_dataset(dataset_path)?;
    run_dataset(conn, &dataset, k)
}

pub fn run_dataset(
    conn: &Connection,
    dataset: &GoldenDataset,
    k: usize,
) -> Result<GoldenEvalReport> {
    if dataset.has_fixture_corpus() {
        evaluate_dataset_with_fixture_corpus(dataset, k)
    } else {
        evaluate_dataset(conn, dataset, k)
    }
}

pub fn evaluate_dataset_with_fixture_corpus(
    dataset: &GoldenDataset,
    k: usize,
) -> Result<GoldenEvalReport> {
    validate_dataset(dataset)?;
    let conn = Connection::open_in_memory().context("open in-memory golden eval DB")?;
    crate::migrate::run_migrations(&conn).context("migrate in-memory golden eval DB")?;
    seed_fixture_corpus(&conn, &dataset.corpus)?;
    evaluate_dataset(&conn, dataset, k)
}

pub fn evaluate_dataset(
    conn: &Connection,
    dataset: &GoldenDataset,
    k: usize,
) -> Result<GoldenEvalReport> {
    validate_dataset(dataset)?;
    let k = k.max(1);
    let fetch_limit = k.max(RANK_K) as i64;
    let mut query_reports = Vec::with_capacity(dataset.queries.len());
    let mut overall_sums = MetricSums::default();
    let mut categories = BTreeMap::<String, CategoryAccumulator>::new();
    let mut slices = BTreeMap::<String, CategoryAccumulator>::new();
    let mut skipped_queries = 0usize;
    let mut abstention_queries = 0usize;
    let mut abstention_passed = 0usize;

    for query in &dataset.queries {
        let results = crate::retrieval::search::search_with_branch(
            conn,
            Some(&query.query),
            query.project.as_deref(),
            query.memory_type.as_deref(),
            fetch_limit,
            0,
            false,
            query.branch.as_deref(),
        )?;
        let evaluation = evaluate_query(query, &results, k);
        let category = categories.entry(query.category.clone()).or_default();
        record_bucket(category, query, &evaluation);
        let slice = slices.entry(query.slice_label().to_string()).or_default();
        record_bucket(slice, query, &evaluation);

        if query.expects_abstention() {
            abstention_queries += 1;
            if evaluation.status == QueryStatus::Pass {
                abstention_passed += 1;
            }
        } else if let Some(metrics) = evaluation.metrics.as_ref() {
            overall_sums.add(metrics);
        } else {
            skipped_queries += 1;
        }

        query_reports.push(evaluation);
    }

    Ok(GoldenEvalReport {
        evaluation_layers: EvaluationLayers::deterministic_retrieval_only(),
        version: dataset.version.clone(),
        description: dataset.description.clone(),
        k,
        rank_k: RANK_K,
        total_queries: dataset.queries.len(),
        scored_queries: overall_sums.averages().map_or(0, |metrics| metrics.count),
        skipped_queries,
        abstention_queries,
        abstention_passed,
        overall: overall_sums.averages(),
        by_slice: slices
            .into_iter()
            .map(|(name, slice)| (name, bucket_evaluation(slice)))
            .collect(),
        by_category: categories
            .into_iter()
            .map(|(name, category)| (name, bucket_evaluation(category)))
            .collect(),
        queries: query_reports,
    })
}

fn record_bucket(
    bucket: &mut CategoryAccumulator,
    query: &GoldenQuery,
    evaluation: &QueryEvaluation,
) {
    bucket.total_queries += 1;
    if query.expects_abstention() {
        bucket.abstention_queries += 1;
        if evaluation.status == QueryStatus::Pass {
            bucket.abstention_passed += 1;
        }
    } else if let Some(metrics) = evaluation.metrics.as_ref() {
        bucket.metrics.add(metrics);
    }
}

fn bucket_evaluation(bucket: CategoryAccumulator) -> CategoryEvaluation {
    CategoryEvaluation {
        total_queries: bucket.total_queries,
        scored_queries: bucket.metrics.averages().map_or(0, |m| m.count),
        abstention_queries: bucket.abstention_queries,
        abstention_passed: bucket.abstention_passed,
        metrics: bucket.metrics.averages(),
    }
}

fn validate_dataset(dataset: &GoldenDataset) -> Result<()> {
    if dataset.queries.is_empty() {
        return Err(anyhow!(
            "golden eval dataset must contain at least one query"
        ));
    }
    validate_fixture_corpus(&dataset.corpus)?;
    let mut seen_ids = HashSet::new();
    for query in &dataset.queries {
        if query.id.trim().is_empty() {
            return Err(anyhow!("golden eval query id must not be empty"));
        }
        if !seen_ids.insert(query.id.as_str()) {
            return Err(anyhow!("duplicate golden eval query id {}", query.id));
        }
        if query.query.trim().is_empty() {
            return Err(anyhow!(
                "golden eval query {} text must not be empty",
                query.id
            ));
        }
        if query.category.trim().is_empty() {
            return Err(anyhow!(
                "golden eval query {} category must not be empty",
                query.id
            ));
        }
        if query
            .slice
            .as_deref()
            .is_some_and(|slice| slice.trim().is_empty())
        {
            return Err(anyhow!(
                "golden eval query {} slice must not be empty",
                query.id
            ));
        }
        if query.expects_abstention()
            && (!query.evidence_refs.is_empty() || !query.relevant_ids.is_empty())
        {
            return Err(anyhow!(
                "golden eval query {} abstention case must not declare expected evidence",
                query.id
            ));
        }
        for evidence_ref in &query.evidence_refs {
            if !evidence_ref.has_match_criteria() {
                return Err(anyhow!(
                    "golden eval query {} contains an empty evidence ref",
                    query.id
                ));
            }
        }
    }
    validate_expected_refs_backed_by_corpus(dataset)?;
    Ok(())
}

fn validate_fixture_corpus(corpus: &[GoldenMemory]) -> Result<()> {
    let mut seen_topic_keys = HashSet::new();
    for (index, memory) in corpus.iter().enumerate() {
        let label = corpus_memory_label(index, memory);
        if memory.project.trim().is_empty() {
            return Err(anyhow!(
                "golden eval corpus memory {label} project must not be empty"
            ));
        }
        if memory.title.trim().is_empty() {
            return Err(anyhow!(
                "golden eval corpus memory {label} title must not be empty"
            ));
        }
        if memory.content.trim().is_empty() {
            return Err(anyhow!(
                "golden eval corpus memory {label} content must not be empty"
            ));
        }
        if memory.memory_type.trim().is_empty() {
            return Err(anyhow!(
                "golden eval corpus memory {label} memory_type must not be empty"
            ));
        }
        if memory.scope.trim().is_empty() {
            return Err(anyhow!(
                "golden eval corpus memory {label} scope must not be empty"
            ));
        }
        if !matches!(memory.status.as_str(), "active" | "stale" | "archived") {
            return Err(anyhow!(
                "golden eval corpus memory {label} status must be active, stale, or archived"
            ));
        }
        if memory
            .branch
            .as_deref()
            .is_some_and(|branch| branch.trim().is_empty())
        {
            return Err(anyhow!(
                "golden eval corpus memory {label} branch must not be empty"
            ));
        }
        if let Some(topic_key) = memory.topic_key.as_deref() {
            if topic_key.trim().is_empty() {
                return Err(anyhow!(
                    "golden eval corpus memory {label} topic_key must not be empty"
                ));
            }
            let key = (
                memory.project.clone(),
                memory.scope.clone(),
                topic_key.to_string(),
            );
            if !seen_topic_keys.insert(key) {
                return Err(anyhow!(
                    "duplicate golden eval corpus topic_key {} for project {} scope {}",
                    topic_key,
                    memory.project,
                    memory.scope
                ));
            }
        }
    }
    Ok(())
}

fn validate_expected_refs_backed_by_corpus(dataset: &GoldenDataset) -> Result<()> {
    if dataset.corpus.is_empty() {
        return Ok(());
    }

    let fixture_memories: Vec<_> = dataset
        .corpus
        .iter()
        .enumerate()
        .map(|(index, memory)| fixture_memory_for_validation(index, memory))
        .collect();

    for query in &dataset.queries {
        if query.expects_abstention() {
            continue;
        }
        for expected_ref in query.expected_refs() {
            let backed = fixture_memories.iter().any(|memory| {
                corpus_memory_matches_query_filter(memory, query) && expected_ref.matches(memory)
            });
            if !backed {
                return Err(anyhow!(
                    "golden eval query {} expected evidence ref is not backed by fixture corpus: {:?}",
                    query.id,
                    expected_ref
                ));
            }
        }
    }
    Ok(())
}

fn fixture_memory_for_validation(index: usize, memory: &GoldenMemory) -> crate::memory::Memory {
    crate::memory::Memory {
        id: index as i64 + 1,
        session_id: Some(format!("golden-eval-{}", index + 1)),
        project: memory.project.clone(),
        topic_key: memory.topic_key.clone(),
        title: memory.title.clone(),
        text: memory.content.clone(),
        memory_type: memory.memory_type.clone(),
        files: memory.files.clone(),
        created_at_epoch: memory.created_at_epoch.unwrap_or(0),
        updated_at_epoch: memory.created_at_epoch.unwrap_or(0),
        status: memory.status.clone(),
        branch: memory.branch.clone(),
        scope: memory.scope.clone(),
    }
}

fn corpus_memory_matches_query_filter(memory: &crate::memory::Memory, query: &GoldenQuery) -> bool {
    if memory.status != "active" {
        return false;
    }
    if let Some(project) = query.project.as_deref() {
        if !crate::project_id::project_matches(Some(&memory.project), project) {
            return false;
        }
    }
    if let Some(branch) = query.branch.as_deref() {
        if memory.branch.as_deref() != Some(branch) {
            return false;
        }
    }
    if let Some(memory_type) = query.memory_type.as_deref() {
        if memory.memory_type != memory_type {
            return false;
        }
    }
    true
}

fn seed_fixture_corpus(conn: &Connection, corpus: &[GoldenMemory]) -> Result<()> {
    for (index, memory) in corpus.iter().enumerate() {
        let id = crate::memory::insert_memory_full(
            conn,
            Some("golden-eval"),
            &memory.project,
            memory.topic_key.as_deref(),
            &memory.title,
            &memory.content,
            &memory.memory_type,
            memory.files.as_deref(),
            memory.branch.as_deref(),
            &memory.scope,
            memory.created_at_epoch,
        )
        .with_context(|| {
            format!(
                "seed golden eval corpus memory {}",
                corpus_memory_label(index, memory)
            )
        })?;
        if memory.status != "active" {
            conn.execute(
                "UPDATE memories SET status = ?1 WHERE id = ?2",
                params![memory.status, id],
            )
            .with_context(|| {
                format!(
                    "set golden eval corpus memory {} status",
                    corpus_memory_label(index, memory)
                )
            })?;
        }
    }
    Ok(())
}

fn corpus_memory_label(index: usize, memory: &GoldenMemory) -> String {
    memory
        .topic_key
        .as_deref()
        .map(str::to_string)
        .unwrap_or_else(|| format!("#{}", index + 1))
}

fn evaluate_query(
    query: &GoldenQuery,
    results: &[crate::memory::Memory],
    k: usize,
) -> QueryEvaluation {
    let expected_refs = query.expected_refs();
    if query.expects_abstention() {
        return QueryEvaluation {
            id: query.id.clone(),
            query: query.query.clone(),
            category: query.category.clone(),
            slice: query.slice_label().to_string(),
            status: if results.is_empty() {
                QueryStatus::Pass
            } else {
                QueryStatus::Fail
            },
            result_count: results.len(),
            retrieved_ids: retrieved_ids(results, k),
            expected_relevant_ids: expected_relevant_ids(&expected_refs),
            missing_relevant_ids: Vec::new(),
            missing_evidence_refs: Vec::new(),
            matched_refs: 0,
            expected_refs: expected_refs.len(),
            metrics: None,
        };
    }

    if expected_refs.is_empty() {
        return QueryEvaluation {
            id: query.id.clone(),
            query: query.query.clone(),
            category: query.category.clone(),
            slice: query.slice_label().to_string(),
            status: QueryStatus::Skip,
            result_count: results.len(),
            retrieved_ids: retrieved_ids(results, k),
            expected_relevant_ids: Vec::new(),
            missing_relevant_ids: Vec::new(),
            missing_evidence_refs: Vec::new(),
            matched_refs: 0,
            expected_refs: 0,
            metrics: None,
        };
    }

    let metrics = score_results(results, &expected_refs, k);
    let matched_ref_indexes = matched_ref_indexes(results, &expected_refs, k);
    let missing_evidence_refs = missing_evidence_refs(&expected_refs, &matched_ref_indexes);
    let missing_relevant_ids = missing_relevant_ids(&missing_evidence_refs);
    let matched_refs = matched_ref_indexes.len();
    let status = if metrics.hit_at_k > 0.0 {
        QueryStatus::Hit
    } else {
        QueryStatus::Miss
    };
    QueryEvaluation {
        id: query.id.clone(),
        query: query.query.clone(),
        category: query.category.clone(),
        slice: query.slice_label().to_string(),
        status,
        result_count: results.len(),
        retrieved_ids: retrieved_ids(results, k),
        expected_relevant_ids: expected_relevant_ids(&expected_refs),
        missing_relevant_ids,
        missing_evidence_refs,
        matched_refs,
        expected_refs: expected_refs.len(),
        metrics: Some(metrics),
    }
}

fn score_results(
    results: &[crate::memory::Memory],
    expected_refs: &[EvidenceRef],
    k: usize,
) -> QueryMetrics {
    let top_k = k.min(results.len());
    let top_rank_k = RANK_K.min(results.len());
    let relevance_at_k: Vec<bool> = results
        .iter()
        .take(top_k)
        .map(|memory| {
            expected_refs
                .iter()
                .any(|evidence_ref| evidence_ref.matches(memory))
        })
        .collect();
    let relevance_at_rank_k: Vec<bool> = results
        .iter()
        .take(top_rank_k)
        .map(|memory| {
            expected_refs
                .iter()
                .any(|evidence_ref| evidence_ref.matches(memory))
        })
        .collect();
    let matched_refs = matched_ref_indexes(results, expected_refs, k).len();
    let relevant_hits = relevance_at_k
        .iter()
        .filter(|is_relevant| **is_relevant)
        .count();
    let precision_denominator = top_k.max(1);

    QueryMetrics {
        hit_at_k: if relevant_hits > 0 { 1.0 } else { 0.0 },
        mrr_at_10: reciprocal_rank_from_relevance(&relevance_at_rank_k),
        precision_at_k: relevant_hits as f64 / precision_denominator as f64,
        recall_at_k: matched_refs as f64 / expected_refs.len() as f64,
        ndcg_at_10: ndcg_at_k(results, expected_refs, RANK_K),
        evidence_recall_at_k: matched_refs as f64 / expected_refs.len() as f64,
    }
}

fn matched_ref_indexes(
    results: &[crate::memory::Memory],
    expected_refs: &[EvidenceRef],
    k: usize,
) -> HashSet<usize> {
    let mut matched = HashSet::new();
    for memory in results.iter().take(k) {
        for (index, evidence_ref) in expected_refs.iter().enumerate() {
            if evidence_ref.matches(memory) {
                matched.insert(index);
            }
        }
    }
    matched
}

fn retrieved_ids(results: &[crate::memory::Memory], k: usize) -> Vec<i64> {
    results.iter().take(k).map(|memory| memory.id).collect()
}

fn expected_relevant_ids(expected_refs: &[EvidenceRef]) -> Vec<i64> {
    expected_refs
        .iter()
        .filter_map(|evidence_ref| evidence_ref.memory_id)
        .collect()
}

fn missing_evidence_refs(
    expected_refs: &[EvidenceRef],
    matched_ref_indexes: &HashSet<usize>,
) -> Vec<EvidenceRef> {
    expected_refs
        .iter()
        .enumerate()
        .filter_map(|(index, evidence_ref)| {
            (!matched_ref_indexes.contains(&index)).then_some(evidence_ref.clone())
        })
        .collect()
}

fn missing_relevant_ids(missing_evidence_refs: &[EvidenceRef]) -> Vec<i64> {
    missing_evidence_refs
        .iter()
        .filter_map(|evidence_ref| evidence_ref.memory_id)
        .collect()
}

fn reciprocal_rank_from_relevance(relevance: &[bool]) -> f64 {
    relevance
        .iter()
        .position(|is_relevant| *is_relevant)
        .map_or(0.0, |index| 1.0 / (index as f64 + 1.0))
}

fn ndcg_at_k(results: &[crate::memory::Memory], expected_refs: &[EvidenceRef], k: usize) -> f64 {
    if k == 0 || expected_refs.is_empty() {
        return 0.0;
    }

    let matches_by_rank: Vec<Vec<usize>> = results
        .iter()
        .take(k)
        .map(|memory| {
            expected_refs
                .iter()
                .enumerate()
                .filter_map(|(index, evidence_ref)| evidence_ref.matches(memory).then_some(index))
                .collect()
        })
        .collect();
    let dcg = best_dcg_for_matchable_ranks(&matches_by_rank);
    let ideal_hits = expected_refs.len().min(k);
    let idcg: f64 = (0..ideal_hits)
        .map(|index| 1.0 / (index as f64 + 2.0).log2())
        .sum();
    if idcg == 0.0 {
        0.0
    } else {
        dcg / idcg
    }
}

fn best_dcg_for_matchable_ranks(matches_by_rank: &[Vec<usize>]) -> f64 {
    let mut best = 0.0;
    for mask in 1usize..(1usize << matches_by_rank.len()) {
        let dcg: f64 = (0..matches_by_rank.len())
            .filter(|rank| (mask & (1usize << rank)) != 0)
            .map(|rank| 1.0 / (rank as f64 + 2.0).log2())
            .sum();
        if dcg > best && can_assign_unique_refs(matches_by_rank, mask) {
            best = dcg;
        }
    }
    best
}

fn can_assign_unique_refs(matches_by_rank: &[Vec<usize>], mask: usize) -> bool {
    let mut selected_ranks: Vec<usize> = (0..matches_by_rank.len())
        .filter(|rank| (mask & (1usize << rank)) != 0)
        .collect();
    if selected_ranks
        .iter()
        .any(|rank| matches_by_rank[*rank].is_empty())
    {
        return false;
    }
    selected_ranks.sort_by_key(|rank| matches_by_rank[*rank].len());

    let mut assigned_refs = HashSet::new();
    assign_rank_ref(0, &selected_ranks, matches_by_rank, &mut assigned_refs)
}

fn assign_rank_ref(
    selected_index: usize,
    selected_ranks: &[usize],
    matches_by_rank: &[Vec<usize>],
    assigned_refs: &mut HashSet<usize>,
) -> bool {
    let Some(rank) = selected_ranks.get(selected_index).copied() else {
        return true;
    };

    for ref_index in &matches_by_rank[rank] {
        if assigned_refs.insert(*ref_index) {
            if assign_rank_ref(
                selected_index + 1,
                selected_ranks,
                matches_by_rank,
                assigned_refs,
            ) {
                return true;
            }
            assigned_refs.remove(ref_index);
        }
    }
    false
}
