use std::sync::Arc;

use tantivy::collector::{Collector, SegmentCollector};
use tantivy::columnar::ColumnValues;
use tantivy::query::{EnableScoring, Query};
use tantivy::schema::Field;
use tantivy::{DocId, DocSet, Score, Searcher, SegmentOrdinal, SegmentReader, TERMINATED};

use crate::TidalError;
use crate::schema::EntityId;

// ---- AllScoresCollector --------------------------------------------------------

/// Returns ALL matching documents with their BM25 scores.
///
/// Unlike `TopDocs`, this collector does not truncate results. Every document
/// matching the query is returned as an `(EntityId, f32)` pair. This is the
/// building block for `TidalDB`'s hybrid ranking pipeline, where BM25 scores are
/// one signal among many and the final top-K selection is done by the ranker,
/// not the text index.
///
/// `requires_scoring()` returns `true`; without this, Tantivy skips BM25
/// computation and every document receives a score of 0.0.
pub struct AllScoresCollector {
    /// The Tantivy `Field` handle for `entity_id` (u64, FAST).
    pub entity_id_field: Field,
}

/// Per-segment collector that reads the `entity_id` fast field and accumulates
/// `(EntityId, f32)` pairs.
///
/// Created by [`AllScoresCollector::for_segment`]; not intended for direct
/// construction by callers.
pub struct AllScoresSegmentCollector {
    /// Fast-field column for `entity_id` values.
    entity_id_col: Arc<dyn ColumnValues<u64>>,
    /// Accumulated results for this segment.
    results: Vec<(EntityId, f32)>,
}

impl Collector for AllScoresCollector {
    type Fruit = Vec<(EntityId, f32)>;
    type Child = AllScoresSegmentCollector;

    fn for_segment(
        &self,
        _segment_ord: SegmentOrdinal,
        reader: &SegmentReader,
    ) -> tantivy::Result<Self::Child> {
        let ff = reader.fast_fields();
        let col = ff.u64("entity_id")?;
        let entity_id_col = col.first_or_default_col(0);
        Ok(AllScoresSegmentCollector {
            entity_id_col,
            results: Vec::new(),
        })
    }

    fn requires_scoring(&self) -> bool {
        true
    }

    fn merge_fruits(
        &self,
        segment_fruits: Vec<Vec<(EntityId, f32)>>,
    ) -> tantivy::Result<Vec<(EntityId, f32)>> {
        let total = segment_fruits.iter().map(Vec::len).sum();
        let mut merged = Vec::with_capacity(total);
        for fruit in segment_fruits {
            merged.extend(fruit);
        }
        Ok(merged)
    }
}

impl SegmentCollector for AllScoresSegmentCollector {
    type Fruit = Vec<(EntityId, f32)>;

    fn collect(&mut self, doc: DocId, score: Score) {
        let eid_val = self.entity_id_col.get_val(doc);
        self.results.push((EntityId::new(eid_val), score));
    }

    fn harvest(self) -> Self::Fruit {
        self.results
    }
}

// ---- score_candidates ----------------------------------------------------------

/// Score a pre-sorted candidate set via BM25 using `DocSet::seek()`.
///
/// This is the targeted scoring path: given a set of entities that have already
/// been selected by another pipeline stage (e.g. vector search), compute their
/// BM25 scores without scanning the full inverted index.
///
/// # Arguments
///
/// * `searcher` - Tantivy searcher for the current reader snapshot.
/// * `query` - The parsed text query whose BM25 scores we want.
/// * `candidates` - Triples of `(segment_ord, doc_id, entity_id)`, **sorted
///   ascending by `(segment_ord, doc_id)`**. The sort order is required because
///   `DocSet::seek()` only moves forward.
///
/// # Errors
///
/// Returns `TidalError::Internal` if Tantivy fails to create a weight or scorer.
pub fn score_candidates(
    searcher: &Searcher,
    query: &dyn Query,
    candidates: &[(u32, u32, EntityId)],
) -> crate::Result<Vec<(EntityId, f32)>> {
    if candidates.is_empty() {
        return Ok(Vec::new());
    }

    // Build a BM25 weight from the query. This requires full scoring statistics.
    let weight = query
        .weight(EnableScoring::enabled_from_searcher(searcher))
        .map_err(|e| {
            TidalError::internal("text_score_candidates", format!("tantivy weight: {e}"))
        })?;

    let mut results = Vec::with_capacity(candidates.len());

    // Group candidates by segment_ord to reuse scorers.
    let mut i = 0;
    while i < candidates.len() {
        let seg_ord = candidates[i].0;

        // Find the extent of candidates in this segment.
        let seg_start = i;
        while i < candidates.len() && candidates[i].0 == seg_ord {
            i += 1;
        }
        let seg_candidates = &candidates[seg_start..i];

        // Obtain the segment reader and create a scorer.
        let segment_reader = searcher.segment_reader(seg_ord);
        let Ok(mut scorer) = weight.scorer(segment_reader, 1.0) else {
            // No postings for this query in this segment -- skip all candidates.
            continue;
        };

        // The scorer is positioned at its first matching doc after creation.
        // If the query has no matches in this segment, the scorer starts at
        // TERMINATED -- skip immediately.
        if scorer.doc() == TERMINATED {
            continue;
        }

        // `seek(target)` moves forward to `target` or the next doc >= target.
        // The `debug_assert!(self.doc() <= target)` in Tantivy's default seek
        // means we must not seek backwards; candidates must be sorted ascending.
        for &(_, doc_id, entity_id) in seg_candidates {
            // If the scorer has already advanced past our candidate, skip.
            if scorer.doc() > doc_id {
                continue;
            }
            let reached = scorer.seek(doc_id);
            if reached == doc_id {
                results.push((entity_id, scorer.score()));
            }
            // If TERMINATED, all remaining candidates in this segment will miss.
            if reached == TERMINATED {
                break;
            }
        }
    }

    Ok(results)
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use std::collections::HashMap;

    use super::*;
    use crate::schema::{TextFieldDef, TextFieldType};
    use crate::text::TextIndex;

    use tantivy::query::QueryParser;

    fn title_field_defs() -> Vec<TextFieldDef> {
        vec![TextFieldDef {
            key: "title".to_owned(),
            field_type: TextFieldType::Text,
        }]
    }

    fn make_metadata(pairs: &[(&str, &str)]) -> HashMap<String, String> {
        pairs
            .iter()
            .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
            .collect()
    }

    /// Helper: create an ephemeral index, insert docs, commit, reload, rebuild entity map.
    fn setup_index(docs: &[(u64, &str)]) -> TextIndex {
        let idx = TextIndex::ephemeral(&title_field_defs()).unwrap();
        {
            let mut w = idx.writer_guard().unwrap();
            for &(eid, title) in docs {
                w.index_item(EntityId::new(eid), &make_metadata(&[("title", title)]))
                    .unwrap();
            }
            w.commit(1).unwrap();
        }
        idx.reader.reload().unwrap();
        idx.rebuild_entity_map().unwrap();
        idx
    }

    #[test]
    fn all_scores_requires_scoring() {
        let idx = TextIndex::ephemeral(&[]).unwrap();
        let c = AllScoresCollector {
            entity_id_field: idx.fields().entity_id,
        };
        assert!(c.requires_scoring());
        idx.close().unwrap();
    }

    #[test]
    fn all_scores_collector_captures_bm25() {
        let idx = setup_index(&[(1, "jazz piano"), (2, "rock guitar"), (3, "jazz violin")]);

        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("jazz").unwrap();

        let collector = AllScoresCollector {
            entity_id_field: idx.fields().entity_id,
        };
        let results = searcher.search(&query, &collector).unwrap();

        // Should find entities 1 and 3, not 2.
        let mut found_ids: Vec<u64> = results.iter().map(|(eid, _)| eid.as_u64()).collect();
        found_ids.sort_unstable();
        assert_eq!(found_ids, vec![1, 3], "should match entities 1 and 3");

        // All scores must be positive.
        for (_eid, score) in &results {
            assert!(*score > 0.0, "BM25 score should be > 0, got {score}");
        }

        idx.close().unwrap();
    }

    #[test]
    fn all_scores_empty_query_returns_empty() {
        let idx = setup_index(&[(1, "jazz piano")]);

        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("nonexistent").unwrap();

        let collector = AllScoresCollector {
            entity_id_field: idx.fields().entity_id,
        };
        let results = searcher.search(&query, &collector).unwrap();
        assert!(results.is_empty(), "no documents should match");

        idx.close().unwrap();
    }

    #[test]
    fn entity_map_rebuilds_after_commit() {
        let idx = setup_index(&[(1, "alpha"), (2, "beta")]);

        assert!(
            idx.entity_doc_address(EntityId::new(1)).is_some(),
            "entity 1 should be in map"
        );
        assert!(
            idx.entity_doc_address(EntityId::new(2)).is_some(),
            "entity 2 should be in map"
        );
        assert!(
            idx.entity_doc_address(EntityId::new(99)).is_none(),
            "entity 99 should not be in map"
        );

        idx.close().unwrap();
    }

    #[test]
    fn entity_map_is_cleared_on_rebuild() {
        let idx = TextIndex::ephemeral(&title_field_defs()).unwrap();

        // Index doc 1, commit, rebuild.
        {
            let mut w = idx.writer_guard().unwrap();
            w.index_item(EntityId::new(1), &make_metadata(&[("title", "first")]))
                .unwrap();
            w.commit(1).unwrap();
        }
        idx.reader.reload().unwrap();
        idx.rebuild_entity_map().unwrap();
        assert!(idx.entity_doc_address(EntityId::new(1)).is_some());

        // Delete doc 1, add doc 2, commit, rebuild.
        {
            let mut w = idx.writer_guard().unwrap();
            w.delete_item(EntityId::new(1));
            w.index_item(EntityId::new(2), &make_metadata(&[("title", "second")]))
                .unwrap();
            w.commit(2).unwrap();
        }
        idx.reader.reload().unwrap();
        idx.rebuild_entity_map().unwrap();

        assert!(
            idx.entity_doc_address(EntityId::new(1)).is_none(),
            "entity 1 should be gone after delete + rebuild"
        );
        assert!(
            idx.entity_doc_address(EntityId::new(2)).is_some(),
            "entity 2 should be present"
        );

        idx.close().unwrap();
    }

    #[test]
    fn score_candidates_returns_bm25_for_matches() {
        let idx = setup_index(&[(1, "jazz piano"), (2, "rock guitar"), (3, "jazz violin")]);

        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("jazz").unwrap();

        // Build candidate list from entity map (sorted by seg_ord, doc_id).
        let mut candidates = Vec::new();
        for eid_val in [1, 3] {
            if let Some((seg_ord, doc_id)) = idx.entity_doc_address(EntityId::new(eid_val)) {
                candidates.push((seg_ord, doc_id, EntityId::new(eid_val)));
            }
        }
        candidates.sort_by_key(|&(s, d, _)| (s, d));

        let results = score_candidates(&searcher, query.as_ref(), &candidates).unwrap();

        assert_eq!(results.len(), 2, "both candidates should score");
        for (eid, score) in &results {
            assert!(
                *score > 0.0,
                "entity {} score should be > 0, got {score}",
                eid.as_u64()
            );
        }

        idx.close().unwrap();
    }

    #[test]
    fn score_candidates_missing_candidate_skipped() {
        let idx = setup_index(&[(1, "jazz piano")]);

        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("jazz").unwrap();

        // Candidate entity 999 does not exist in the index at all.
        // We fabricate a plausible (seg_ord=0, doc_id=999) for it.
        let candidates = vec![(0, 999, EntityId::new(999))];

        let results = score_candidates(&searcher, query.as_ref(), &candidates).unwrap();
        let found_ids: Vec<u64> = results.iter().map(|(eid, _)| eid.as_u64()).collect();
        assert!(
            !found_ids.contains(&999),
            "entity 999 should not appear in results"
        );
    }

    #[test]
    fn score_candidates_empty_candidates() {
        let idx = setup_index(&[(1, "jazz piano")]);

        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("jazz").unwrap();

        let results = score_candidates(&searcher, query.as_ref(), &[]).unwrap();
        assert!(results.is_empty());
    }

    #[test]
    fn score_candidates_non_matching_query() {
        let idx = setup_index(&[(1, "jazz piano"), (2, "rock guitar")]);

        let searcher = idx.reader.searcher();
        let title_field = idx.fields().text_fields[0].1;
        let qp = QueryParser::for_index(&idx.index, vec![title_field]);
        let query = qp.parse_query("nonexistent").unwrap();

        let mut candidates = Vec::new();
        if let Some((s, d)) = idx.entity_doc_address(EntityId::new(1)) {
            candidates.push((s, d, EntityId::new(1)));
        }

        let results = score_candidates(&searcher, query.as_ref(), &candidates).unwrap();
        assert!(
            results.is_empty(),
            "no candidates should score for unmatched query"
        );
    }
}
