use std::{cmp::Ordering, collections::BinaryHeap};

use tantivy::{
    DocId, Score, SegmentOrdinal, SegmentReader,
    collector::{Collector, SegmentCollector},
    columnar::Column,
    schema::Field,
};

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.
///
/// # Unbounded — prefer [`BoundedScoresCollector`] on the query hot path
///
/// This collector keeps EVERY match, so a broad / near-ubiquitous term over a
/// large corpus materializes one `(EntityId, f32)` per match — an allocation
/// proportional to corpus match count, not the requested limit. The SEARCH
/// pipeline uses [`BoundedScoresCollector`] instead, which caps the result to a
/// top-K inside Tantivy. `AllScoresCollector` is retained for callers (and tests)
/// that genuinely need the full match set.
pub struct AllScoresCollector {
    /// The Tantivy `Field` handle for `entity_id` (u64, FAST).
    ///
    /// [`for_segment`](Collector::for_segment) resolves the fast-field column
    /// from this handle (via the segment schema), so the looked-up column is
    /// guaranteed to be the field the caller passed — no hardcoded name.
    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 {
    /// Strict fast-field column for `entity_id` values.
    ///
    /// Held as a `Column<u64>` (not a defaulted `ColumnValues`) so a document
    /// with no `entity_id` value reads as `None` rather than a defaulted `0` —
    /// `0` is a legal [`EntityId`], so a default would silently alias entity 0.
    entity_id_col: Column<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> {
        // Resolve the fast-field column from the *typed* `entity_id_field`
        // handle carried in `self`, not a hardcoded string. The column name is
        // derived from the schema field, so the handle every caller sets and
        // the column we read are provably the same field. The canonical name
        // itself is defined once as `index::ENTITY_ID_FIELD`.
        let column_name = reader.schema().get_field_name(self.entity_id_field);
        let entity_id_col = reader.fast_fields().u64(column_name)?;
        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) {
        // `first` is the strict accessor: `None` means this doc carries no
        // `entity_id` value at all. Every document written by the index path
        // (`writer::index_item`) sets one, so `None` is a should-be-impossible
        // state — we skip it rather than fabricate a result for entity 0, which
        // is a legal id and would silently corrupt the result set.
        if let Some(eid_val) = self.entity_id_col.first(doc) {
            self.results.push((EntityId::new(eid_val), score));
        }
    }

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

// ---- BoundedScoresCollector ----------------------------------------------------

/// Like [`AllScoresCollector`] but keeps only the `cap` highest-scoring matches.
///
/// On a broad or near-ubiquitous term over a large corpus, [`AllScoresCollector`]
/// allocates one `(EntityId, f32)` per match (plus a downstream `HashMap`) — a
/// memory-amplification / latency-spike vector proportional to corpus match count
/// rather than the requested limit. This collector keeps a per-segment bounded
/// min-heap of size `cap`, so peak memory is `O(cap)` INSIDE Tantivy regardless of
/// how many documents match. A `cap` of `0` collects nothing.
///
/// The returned set is the global top-`cap` by score; `NaN` scores (degenerate
/// BM25 term statistics) are treated as the smallest and evicted first, matching
/// the pipeline's `finite` neutralization. Ordering of the returned `Vec` is
/// unspecified — the SEARCH pipeline sorts it deterministically (descending score,
/// id tie-break) before use.
pub struct BoundedScoresCollector {
    /// The Tantivy `Field` handle for `entity_id` (u64, FAST). See
    /// [`AllScoresCollector::entity_id_field`].
    pub entity_id_field: Field,
    /// Maximum number of highest-scoring documents to keep.
    pub cap: usize,
}

impl BoundedScoresCollector {
    /// Collector bounded to the `cap` highest-scoring documents.
    #[must_use]
    pub const fn new(entity_id_field: Field, cap: usize) -> Self {
        Self {
            entity_id_field,
            cap,
        }
    }
}

/// A `(score, entity_id)` pair ordered so a [`BinaryHeap`] (a max-heap) behaves as
/// a MIN-heap on score — the smallest score is at the top and is evicted first when
/// the heap is at capacity. `NaN` scores are treated as the smallest (evicted
/// first). Ties break on entity id so eviction is deterministic.
#[derive(Debug, Clone, Copy, PartialEq)]
struct MinScored {
    score: f32,
    entity_id: u64,
}

impl Eq for MinScored {}

impl Ord for MinScored {
    fn cmp(&self, other: &Self) -> Ordering {
        // Reverse the score order so `BinaryHeap::peek`/`pop` yields the SMALLEST
        // score (the eviction candidate). NaN sorts as the smallest score.
        let a = if self.score.is_nan() {
            f32::MIN
        } else {
            self.score
        };
        let b = if other.score.is_nan() {
            f32::MIN
        } else {
            other.score
        };
        b.total_cmp(&a)
            .then_with(|| other.entity_id.cmp(&self.entity_id))
    }
}

impl PartialOrd for MinScored {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// Per-segment bounded collector backing [`BoundedScoresCollector`].
pub struct BoundedScoresSegmentCollector {
    entity_id_col: Column<u64>,
    heap: BinaryHeap<MinScored>,
    cap: usize,
}

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

    fn for_segment(
        &self,
        _segment_ord: SegmentOrdinal,
        reader: &SegmentReader,
    ) -> tantivy::Result<Self::Child> {
        let column_name = reader.schema().get_field_name(self.entity_id_field);
        let entity_id_col = reader.fast_fields().u64(column_name)?;
        Ok(BoundedScoresSegmentCollector {
            entity_id_col,
            heap: BinaryHeap::with_capacity(self.cap.saturating_add(1)),
            cap: self.cap,
        })
    }

    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);
        }
        // Each segment already trimmed to `cap`, but the UNION of per-segment
        // top-`cap` sets can exceed `cap`. Trim to the GLOBAL top-`cap` so the
        // returned set never exceeds the requested bound.
        if self.cap == 0 {
            return Ok(Vec::new());
        }
        if merged.len() > self.cap {
            // Partition so the `cap` highest scores occupy [0, cap); NaN sinks to
            // the bottom (treated as the smallest), ties broken on ascending id
            // for determinism, then drop the tail.
            let finite = |s: f32| if s.is_nan() { f32::MIN } else { s };
            let nth = self.cap.saturating_sub(1).min(merged.len() - 1);
            merged.select_nth_unstable_by(nth, |a, b| {
                finite(b.1)
                    .total_cmp(&finite(a.1))
                    .then_with(|| a.0.as_u64().cmp(&b.0.as_u64()))
            });
            merged.truncate(self.cap);
        }
        Ok(merged)
    }
}

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

    fn collect(&mut self, doc: DocId, score: Score) {
        if self.cap == 0 {
            return;
        }
        let Some(eid_val) = self.entity_id_col.first(doc) else {
            return;
        };
        self.heap.push(MinScored {
            score,
            entity_id: eid_val,
        });
        // Keep only the top `cap`: once over capacity, evict the smallest score
        // (the heap root). Peak memory stays `O(cap)`.
        if self.heap.len() > self.cap {
            self.heap.pop();
        }
    }

    fn harvest(self) -> Self::Fruit {
        self.heap
            .into_iter()
            .map(|m| (EntityId::new(m.entity_id), m.score))
            .collect()
    }
}

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

    use tantivy::query::QueryParser;

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

    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.
    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().unwrap();
        }
        idx.reader.reload().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_collects_genuine_entity_zero() {
        // Entity 0 is a legal id. The strict `first` accessor must still collect
        // it — the skip path is only for documents with no id value at all.
        let idx = setup_index(&[(0, "jazz piano"), (1, "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();

        let mut found_ids: Vec<u64> = results.iter().map(|(eid, _)| eid.as_u64()).collect();
        found_ids.sort_unstable();
        assert_eq!(found_ids, vec![0, 1], "genuine entity 0 must be collected");

        idx.close().unwrap();
    }

    #[test]
    fn all_scores_resolves_column_from_typed_handle() {
        // The collector must read whatever column the `entity_id_field` handle
        // names — proving `for_segment` resolves the column from the handle, not
        // a hardcoded string. The schema names it `index::ENTITY_ID_FIELD`.
        let idx = setup_index(&[(42, "jazz piano")]);

        let searcher = idx.reader.searcher();
        let resolved = searcher
            .schema()
            .get_field_name(idx.fields().entity_id)
            .to_owned();
        assert_eq!(
            resolved,
            crate::text::index::ENTITY_ID_FIELD,
            "entity_id handle must name the canonical column"
        );

        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();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0.as_u64(), 42);

        idx.close().unwrap();
    }

    /// W regression: a bounded collector must return at most `cap` documents even
    /// when more match, and must keep the highest-scoring ones. Doc 1 repeats the
    /// query term so its BM25 score outranks docs 2 and 3; `cap = 1` must keep
    /// exactly doc 1.
    #[test]
    fn bounded_collector_keeps_top_cap_by_score() {
        // Doc 1 mentions "jazz" twice → higher BM25 than docs 2/3 (once each).
        let idx = setup_index(&[
            (1, "jazz jazz piano"),
            (2, "jazz 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();

        // cap = 1: only the single highest-scoring match (doc 1) survives.
        let bounded = BoundedScoresCollector::new(idx.fields().entity_id, 1);
        let results = searcher.search(&query, &bounded).unwrap();
        assert_eq!(
            results.len(),
            1,
            "bounded collector must return at most cap=1"
        );
        assert_eq!(
            results[0].0.as_u64(),
            1,
            "the single retained doc must be the highest-scoring match"
        );

        // cap = 2: the two highest-scoring of three matches survive.
        let bounded2 = BoundedScoresCollector::new(idx.fields().entity_id, 2);
        let results2 = searcher.search(&query, &bounded2).unwrap();
        assert_eq!(
            results2.len(),
            2,
            "bounded collector must return at most cap=2"
        );
        assert!(
            results2.iter().any(|(id, _)| id.as_u64() == 1),
            "the highest-scoring doc must be retained"
        );

        // The unbounded collector still returns all three matches.
        let unbounded = AllScoresCollector {
            entity_id_field: idx.fields().entity_id,
        };
        let all = searcher.search(&query, &unbounded).unwrap();
        assert_eq!(all.len(), 3, "unbounded collector returns every match");

        idx.close().unwrap();
    }

    /// A `cap` of 0 collects nothing.
    #[test]
    fn bounded_collector_cap_zero_collects_nothing() {
        let idx = setup_index(&[(1, "jazz piano"), (2, "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 bounded = BoundedScoresCollector::new(idx.fields().entity_id, 0);
        let results = searcher.search(&query, &bounded).unwrap();
        assert!(results.is_empty(), "cap=0 must collect nothing");
        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();
    }
}
