//! Sorted term dictionary with binary search lookup.
//!
//! Maps terms (strings) to posting list offsets. Terms must be added in
//! sorted order. On-disk format:
//!
//! ```text
//! [num_terms: u32]
//! [term_data: [len: u16] [bytes] [offset: u64] ...]
//! [offset_index: [u32 per term]]  // byte position of each term entry in term_data
//! ```
//!
//! The offset index enables O(log n) binary search without scanning all term
//! bytes sequentially.
//!
//! See [[inverted-index]] and [[architecture-overview#Step 2]].

/// Builds a term dictionary. Terms must be added in lexicographic order.
pub struct TermDictBuilder {
    /// Serialized term entries: [len: u16][bytes][offset: u64]
    term_data: Vec<u8>,
    /// Byte position in term_data where each term entry starts.
    entry_offsets: Vec<u32>,
    last_term: Option<String>,
}

impl TermDictBuilder {
    pub fn new() -> Self {
        Self {
            term_data: Vec::new(),
            entry_offsets: Vec::new(),
            last_term: None,
        }
    }

    /// Add a term → offset mapping. Terms must be added in sorted order.
    pub fn add(&mut self, term: &str, offset: u64) {
        if let Some(ref last) = self.last_term {
            debug_assert!(
                term > last.as_str(),
                "terms must be added in sorted order: {term:?} <= {last:?}"
            );
        }

        let entry_start = self.term_data.len() as u32;
        self.entry_offsets.push(entry_start);

        let bytes = term.as_bytes();
        let len = bytes.len() as u16;
        self.term_data.extend_from_slice(&len.to_le_bytes());
        self.term_data.extend_from_slice(bytes);
        self.term_data.extend_from_slice(&offset.to_le_bytes());

        self.last_term = Some(term.to_string());
    }

    /// Finalize and return the encoded term dictionary bytes.
    pub fn finish(self) -> Vec<u8> {
        let num_terms = self.entry_offsets.len() as u32;
        let mut result = Vec::new();

        // Header: num_terms
        result.extend_from_slice(&num_terms.to_le_bytes());

        // Term data
        result.extend_from_slice(&self.term_data);

        // Offset index (one u32 per term)
        for &off in &self.entry_offsets {
            result.extend_from_slice(&off.to_le_bytes());
        }

        result
    }
}

impl Default for TermDictBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Read-only term dictionary with O(log n) binary search.
pub struct TermDict<'a> {
    data: &'a [u8],
    num_terms: u32,
    /// Byte offset where term_data starts (after the 4-byte header).
    term_data_start: usize,
    /// Byte offset where the offset index starts.
    index_start: usize,
}

impl<'a> TermDict<'a> {
    /// Open a term dictionary from encoded bytes.
    pub fn open(data: &'a [u8]) -> Self {
        if data.len() < 4 {
            return Self {
                data,
                num_terms: 0,
                term_data_start: 4,
                index_start: 4,
            };
        }

        let num_terms = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
        let term_data_start = 4;
        // The offset index is at the end: num_terms * 4 bytes
        let index_size = num_terms as usize * 4;
        let index_start = data.len() - index_size;

        Self {
            data,
            num_terms,
            term_data_start,
            index_start,
        }
    }

    /// Number of terms in the dictionary.
    pub fn len(&self) -> u32 {
        self.num_terms
    }

    /// Whether the dictionary is empty.
    pub fn is_empty(&self) -> bool {
        self.num_terms == 0
    }

    /// Look up a term and return its posting list offset, or None.
    pub fn get(&self, term: &str) -> Option<u64> {
        if self.num_terms == 0 {
            return None;
        }

        let mut lo: u32 = 0;
        let mut hi: u32 = self.num_terms;

        while lo < hi {
            let mid = lo + (hi - lo) / 2;
            let (mid_term, mid_offset) = self.read_entry(mid);

            match mid_term.as_bytes().cmp(term.as_bytes()) {
                std::cmp::Ordering::Equal => return Some(mid_offset),
                std::cmp::Ordering::Less => lo = mid + 1,
                std::cmp::Ordering::Greater => hi = mid,
            }
        }

        None
    }

    /// Iterate all terms starting with `prefix`, yielding (term, offset) pairs.
    ///
    /// Uses binary search to find the first matching term, then scans forward.
    pub fn prefix_iter(&self, prefix: &str) -> PrefixIter<'a, '_> {
        if self.num_terms == 0 {
            return PrefixIter {
                dict: self,
                prefix: prefix.to_string(),
                idx: 0,
                end: 0,
            };
        }

        // Binary search for first term >= prefix
        let mut lo: u32 = 0;
        let mut hi: u32 = self.num_terms;
        while lo < hi {
            let mid = lo + (hi - lo) / 2;
            let (mid_term, _) = self.read_entry(mid);
            if mid_term.as_bytes() < prefix.as_bytes() {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }

        PrefixIter {
            dict: self,
            prefix: prefix.to_string(),
            idx: lo,
            end: self.num_terms,
        }
    }

    /// Read the term entry at the given index. Returns (term_bytes, offset).
    fn read_entry(&self, idx: u32) -> (&'a str, u64) {
        let index_pos = self.index_start + (idx as usize) * 4;
        let entry_offset = u32::from_le_bytes([
            self.data[index_pos],
            self.data[index_pos + 1],
            self.data[index_pos + 2],
            self.data[index_pos + 3],
        ]) as usize;

        let abs_pos = self.term_data_start + entry_offset;
        let len = u16::from_le_bytes([self.data[abs_pos], self.data[abs_pos + 1]]) as usize;
        let term_bytes = &self.data[abs_pos + 2..abs_pos + 2 + len];
        let offset_pos = abs_pos + 2 + len;
        let offset = u64::from_le_bytes([
            self.data[offset_pos],
            self.data[offset_pos + 1],
            self.data[offset_pos + 2],
            self.data[offset_pos + 3],
            self.data[offset_pos + 4],
            self.data[offset_pos + 5],
            self.data[offset_pos + 6],
            self.data[offset_pos + 7],
        ]);

        let term = std::str::from_utf8(term_bytes).expect("term dictionary contains invalid UTF-8");
        (term, offset)
    }
}

/// Iterator over terms matching a prefix.
pub struct PrefixIter<'a, 'b> {
    dict: &'b TermDict<'a>,
    prefix: String,
    idx: u32,
    end: u32,
}

impl<'a, 'b> Iterator for PrefixIter<'a, 'b> {
    type Item = (&'a str, u64);

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx >= self.end {
            return None;
        }
        let (term, offset) = self.dict.read_entry(self.idx);
        if !term.starts_with(&self.prefix) {
            return None; // past the prefix range
        }
        self.idx += 1;
        Some((term, offset))
    }
}

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

    #[test]
    fn empty_dict() {
        let builder = TermDictBuilder::new();
        let data = builder.finish();
        let dict = TermDict::open(&data);
        assert_eq!(dict.len(), 0);
        assert!(dict.is_empty());
        assert_eq!(dict.get("anything"), None);
    }

    #[test]
    fn single_term() {
        let mut builder = TermDictBuilder::new();
        builder.add("hello", 42);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.len(), 1);
        assert_eq!(dict.get("hello"), Some(42));
        assert_eq!(dict.get("world"), None);
    }

    #[test]
    fn multiple_terms() {
        let mut builder = TermDictBuilder::new();
        builder.add("apple", 100);
        builder.add("banana", 200);
        builder.add("cherry", 300);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.len(), 3);
        assert_eq!(dict.get("apple"), Some(100));
        assert_eq!(dict.get("banana"), Some(200));
        assert_eq!(dict.get("cherry"), Some(300));
    }

    #[test]
    fn lookup_miss() {
        let mut builder = TermDictBuilder::new();
        builder.add("apple", 100);
        builder.add("cherry", 300);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.get("banana"), None);
        assert_eq!(dict.get("aaa"), None);
        assert_eq!(dict.get("zzz"), None);
    }

    #[test]
    fn binary_search_first_term() {
        let mut builder = TermDictBuilder::new();
        builder.add("aaa", 1);
        builder.add("bbb", 2);
        builder.add("ccc", 3);
        builder.add("ddd", 4);
        builder.add("eee", 5);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.get("aaa"), Some(1));
    }

    #[test]
    fn binary_search_last_term() {
        let mut builder = TermDictBuilder::new();
        builder.add("aaa", 1);
        builder.add("bbb", 2);
        builder.add("ccc", 3);
        builder.add("ddd", 4);
        builder.add("eee", 5);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.get("eee"), Some(5));
    }

    #[test]
    fn binary_search_middle_term() {
        let mut builder = TermDictBuilder::new();
        builder.add("aaa", 1);
        builder.add("bbb", 2);
        builder.add("ccc", 3);
        builder.add("ddd", 4);
        builder.add("eee", 5);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.get("ccc"), Some(3));
    }

    #[test]
    fn unicode_terms() {
        let mut builder = TermDictBuilder::new();
        // Sorted by byte order
        builder.add("cafe", 1);
        builder.add("caf\u{00e9}", 2); // café
        builder.add("naïve", 3);
        builder.add("über", 4);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.get("cafe"), Some(1));
        assert_eq!(dict.get("caf\u{00e9}"), Some(2));
        assert_eq!(dict.get("naïve"), Some(3));
        assert_eq!(dict.get("über"), Some(4));
    }

    #[test]
    fn large_dict() {
        let count = 10_000;
        let mut builder = TermDictBuilder::new();
        let mut terms: Vec<String> = (0..count).map(|i| format!("term_{i:06}")).collect();
        terms.sort();
        for (i, term) in terms.iter().enumerate() {
            builder.add(term, i as u64);
        }
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.len(), count as u32);

        // Spot-check several lookups
        for (i, term) in terms.iter().enumerate() {
            assert_eq!(dict.get(term), Some(i as u64));
        }
    }

    #[test]
    fn large_offsets() {
        let mut builder = TermDictBuilder::new();
        builder.add("a", u64::MAX - 1);
        builder.add("b", u64::MAX);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.get("a"), Some(u64::MAX - 1));
        assert_eq!(dict.get("b"), Some(u64::MAX));
    }

    #[test]
    fn empty_string_term() {
        let mut builder = TermDictBuilder::new();
        builder.add("", 0);
        builder.add("a", 1);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        assert_eq!(dict.get(""), Some(0));
        assert_eq!(dict.get("a"), Some(1));
    }

    // --- Prefix iteration tests ---

    #[test]
    fn prefix_multiple_matches() {
        let mut builder = TermDictBuilder::new();
        builder.add("apple", 1);
        builder.add("application", 2);
        builder.add("apply", 3);
        builder.add("banana", 4);
        builder.add("cherry", 5);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("app").collect();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0], ("apple", 1));
        assert_eq!(results[1], ("application", 2));
        assert_eq!(results[2], ("apply", 3));
    }

    #[test]
    fn prefix_no_matches() {
        let mut builder = TermDictBuilder::new();
        builder.add("apple", 1);
        builder.add("banana", 2);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("xyz").collect();
        assert!(results.is_empty());
    }

    #[test]
    fn prefix_empty_yields_all() {
        let mut builder = TermDictBuilder::new();
        builder.add("a", 1);
        builder.add("b", 2);
        builder.add("c", 3);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("").collect();
        assert_eq!(results.len(), 3);
    }

    #[test]
    fn prefix_single_char() {
        let mut builder = TermDictBuilder::new();
        builder.add("ax", 1);
        builder.add("ay", 2);
        builder.add("bx", 3);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("a").collect();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].0, "ax");
        assert_eq!(results[1].0, "ay");
    }

    #[test]
    fn prefix_exact_term_match() {
        let mut builder = TermDictBuilder::new();
        builder.add("hello", 1);
        builder.add("helloworld", 2);
        builder.add("help", 3);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("hello").collect();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].0, "hello");
        assert_eq!(results[1].0, "helloworld");
    }

    #[test]
    fn prefix_unicode() {
        let mut builder = TermDictBuilder::new();
        builder.add("cafe", 1);
        builder.add("caf\u{00e9}", 2);
        builder.add("car", 3);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("caf").collect();
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn prefix_on_empty_dict() {
        let builder = TermDictBuilder::new();
        let data = builder.finish();
        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("any").collect();
        assert!(results.is_empty());
    }

    #[test]
    fn prefix_at_end_of_dict() {
        let mut builder = TermDictBuilder::new();
        builder.add("aaa", 1);
        builder.add("zzz", 2);
        let data = builder.finish();

        let dict = TermDict::open(&data);
        let results: Vec<_> = dict.prefix_iter("zz").collect();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, "zzz");
    }
}
