use std::{
    collections::{BTreeMap, btree_map::Entry},
    time::Instant,
};

use anyhow::Result;
use tantivy::TantivyDocument;

use crate::{
    utils::{
        cmp_doc::sort_docs, key::ShardingKey, row_entry::RowEntry, schema::SortField,
        tantivy_doc::TantivyDoc,
    },
    writer::cdc_log::OperationType,
};

pub struct DocStat {
    pub doc_count: usize, // All documents including deleted ones
    pub doc_bytes: usize,
    pub created_time: Instant,
}

#[derive(Default, Debug, Clone)]
pub struct HandleStat {
    pub handle_count: usize,
    pub delete_count: usize,
    pub min_handle: Vec<u8>,
    pub mid_handle: Vec<u8>,
    pub max_handle: Vec<u8>,
}

#[derive(Default)]
pub struct TantivyDocResult {
    pub docs: Vec<TantivyDocument>,
    pub row_entries: Vec<RowEntry>,
    pub handle_stat: HandleStat,
}

// A temporary in-memory buffer that accumulates Tantivy documents
// before they are flushed to the index. The buffer ensures that:
// 1. Only the latest version per handle is retained.
// 2. Deletes are tracked separately for efficiency.
// 3. Lightweight statistics are maintained incrementally.
pub struct TantivyDocBuffer {
    tantivy_docs: BTreeMap<ShardingKey, TantivyDoc>,
    doc_bytes: usize,
    delete_count: usize,
    created_time: Instant,
}

#[allow(clippy::new_without_default)]
impl TantivyDocBuffer {
    pub fn new() -> Self {
        Self {
            tantivy_docs: BTreeMap::new(),
            doc_bytes: 0,
            delete_count: 0,
            created_time: Instant::now(),
        }
    }

    pub fn stat(&self) -> DocStat {
        DocStat {
            doc_count: self.tantivy_docs.len(),
            doc_bytes: self.doc_bytes,
            created_time: self.created_time,
        }
    }

    // Inserts or updates a document in the buffer.
    //
    // - If the handle is new, the doc is inserted.
    // - If the handle exists and the new version is newer, the old doc is replaced.
    // - If the version is not newer, the doc is ignored.
    //
    // Returns `true` if the doc was actually inserted or updated.
    pub fn add(&mut self, doc: TantivyDoc) -> bool {
        match self.tantivy_docs.entry(doc.sharding_key.clone()) {
            Entry::Occupied(mut entry) => {
                if entry.get().version < doc.version {
                    if entry.get().op_type == OperationType::Delete {
                        self.delete_count -= 1;
                    }
                    if doc.op_type == OperationType::Delete {
                        self.delete_count += 1;
                    }
                    self.doc_bytes -= entry.get().doc_bytes;
                    self.doc_bytes += doc.doc_bytes;
                    // Replace the existing document
                    entry.insert(doc);
                    true
                } else {
                    // Ignore if version does not increase
                    false
                }
            }
            Entry::Vacant(entry) => {
                // First time insertion of this handle
                if doc.op_type == OperationType::Delete {
                    self.delete_count += 1;
                }
                self.doc_bytes += doc.doc_bytes;
                entry.insert(doc);
                true
            }
        }
    }

    // Finalizes the buffer into a result set for indexing and deletion.
    //
    // This function consumes the buffer and produces:
    // - A list of Tantivy documents to be added.
    // - A sorted list of handles to be deleted.
    // - Aggregate handle statistics.
    //
    // The buffer itself is cleared after calling this function.
    pub fn finalize(self, columns: &[SortField]) -> Result<TantivyDocResult> {
        if self.tantivy_docs.is_empty() {
            return Ok(TantivyDocResult::default());
        }
        let handle_count = self.tantivy_docs.len();
        let min_handle = self.tantivy_docs.keys().next().unwrap().raw().to_vec();
        let mid_handle = self
            .tantivy_docs
            .iter()
            .nth(handle_count / 2)
            .map(|(h, _)| h.raw().to_vec())
            .unwrap();
        let max_handle = self.tantivy_docs.keys().last().unwrap().raw().to_vec();
        let handle_stat = HandleStat {
            handle_count,
            delete_count: self.delete_count,
            min_handle,
            mid_handle,
            max_handle,
        };

        let mut docs = Vec::with_capacity(self.tantivy_docs.len() - self.delete_count);
        // row_entries is sorted due to BTreeMap iteration order.
        let mut row_entries = Vec::with_capacity(self.tantivy_docs.len());
        for (handle, doc) in self.tantivy_docs {
            let is_deleted = doc.op_type == OperationType::Delete;
            row_entries.push(RowEntry::new(handle.into_vec(), is_deleted, doc.version));
            if !is_deleted {
                docs.push(doc.doc);
            }
        }

        Ok(TantivyDocResult {
            docs: sort_docs(docs, columns),
            row_entries,
            handle_stat,
        })
    }
}

#[cfg(test)]
mod tests {
    use tantivy::TantivyDocument;

    use super::*;
    use crate::{utils::tantivy_doc::TantivyDoc, writer::tantivy_doc_buffer::TantivyDocBuffer};

    fn make_doc(
        handle: &[u8],
        version: u64,
        op_type: OperationType,
        doc_bytes: usize,
    ) -> TantivyDoc {
        TantivyDoc {
            sharding_key: ShardingKey::from_raw(handle.to_vec()),
            version,
            op_type,
            doc: TantivyDocument::default(),
            doc_bytes,
        }
    }

    #[test]
    fn test_insert_and_update_put() {
        let mut buf = TantivyDocBuffer::new();
        let h = b"h1";

        // Insert v1
        assert!(buf.add(make_doc(h, 1, OperationType::Insert, 100)));
        assert_eq!(buf.stat().doc_count, 1);

        // Update with higher version v2
        assert!(buf.add(make_doc(h, 2, OperationType::Update, 200)));
        assert_eq!(buf.stat().doc_bytes, 200);

        // Try adding same version (ignored)
        assert!(!buf.add(make_doc(h, 2, OperationType::Delete, 50)));
    }

    #[test]
    fn test_put_to_delete_transition() {
        let mut buf = TantivyDocBuffer::new();
        let h = b"h1";

        // Insert Put
        assert!(buf.add(make_doc(h, 1, OperationType::Insert, 100)));

        // Replace with Delete
        assert!(buf.add(make_doc(h, 2, OperationType::Delete, 50)));

        // Replace with newer Put (restore)
        assert!(buf.add(make_doc(h, 3, OperationType::Insert, 80)));
    }

    #[test]
    fn test_multiple_handles_stats() {
        let mut buf = TantivyDocBuffer::new();
        buf.add(make_doc(b"a", 1, OperationType::Insert, 100));
        buf.add(make_doc(b"b", 2, OperationType::Insert, 200));
        buf.add(make_doc(b"c", 3, OperationType::Delete, 50));

        let stat = buf.stat();
        assert_eq!(stat.doc_count, 3);

        let result = buf.finalize(&[]).unwrap();
        assert_eq!(result.docs.len(), 2);
        assert_eq!(result.row_entries.len(), 3);
        assert!(!result.row_entries[0].del_mark);
        assert!(!result.row_entries[1].del_mark);
        assert!(result.row_entries[2].del_mark);
        assert_eq!(result.row_entries[0].version, 1);
        assert_eq!(result.row_entries[1].version, 2);
        assert_eq!(result.row_entries[2].version, 3);

        let hs = &result.handle_stat;
        assert_eq!(hs.handle_count, 3);
        assert_eq!(hs.delete_count, 1);
        assert_eq!(hs.min_handle, b"a");
        assert_eq!(hs.max_handle, b"c");
        // middle handle should be "b"
        assert_eq!(hs.mid_handle, b"b");
    }

    #[test]
    fn test_delete_overwrite_same_handle() {
        let mut buf = TantivyDocBuffer::new();
        let h = b"x";

        // Delete v1
        assert!(buf.add(make_doc(h, 1, OperationType::Delete, 10)));

        // Delete again with higher version (should not double count)
        assert!(buf.add(make_doc(h, 2, OperationType::Delete, 15)));
    }

    #[test]
    fn test_finalize_row_entries_keep_latest_version() {
        let mut buf = TantivyDocBuffer::new();
        let h = b"x";

        assert!(buf.add(make_doc(h, 1, OperationType::Insert, 10)));
        assert!(buf.add(make_doc(h, 3, OperationType::Delete, 15)));
        assert!(!buf.add(make_doc(h, 2, OperationType::Insert, 12)));

        let result = buf.finalize(&[]).unwrap();
        assert_eq!(result.row_entries, vec![RowEntry::new(h.to_vec(), true, 3)]);
    }

    #[test]
    fn test_finalize_empty() {
        let buf = TantivyDocBuffer::new();
        let res = buf.finalize(&[]).unwrap();
        assert!(res.docs.is_empty());
        assert!(res.row_entries.is_empty());
        assert_eq!(res.handle_stat.handle_count, 0);
    }

    #[test]
    fn test_version_rollback_protection() {
        let mut buf = TantivyDocBuffer::new();
        let h = b"key1";
        buf.add(make_doc(h, 10, OperationType::Insert, 1000));
        let added = buf.add(make_doc(h, 5, OperationType::Insert, 500));
        assert!(!added);
        assert_eq!(buf.stat().doc_bytes, 1000);
    }

    #[test]
    fn test_mid_handle_even_count() {
        let mut buf = TantivyDocBuffer::new();
        buf.add(make_doc(b"1", 1, OperationType::Insert, 10));
        buf.add(make_doc(b"2", 1, OperationType::Insert, 10));
        let res = buf.finalize(&[]).unwrap();
        assert_eq!(res.handle_stat.mid_handle, b"2");
    }
}