use std::sync::Arc;

use anyhow::{Result, anyhow, bail};
use log::info;
use tici_common::cdc_file_name::CDCFileName;
use tokio::{sync::mpsc, time::Instant};

use crate::{
    ShardWriterMeta,
    config::FragWriterConfig,
    utils::{
        frag_path_generator::FragPathGenerator, metrics::WRITER_FRAG_OPERATION_LATENCY,
        schema::IndexMeta, tantivy_doc::TantivyDoc,
    },
    writer::{
        cdc_log::CDCLogEntry,
        frag_uploader::{FragUploadTask, SubmitInfo},
        frag_writer::{FragWriter, FragWriterBuilder},
    },
};

// CDCLogFragWriter will get new CDCLogEntry from cdc_log_receiver continuously,
// and write them to current fragments.
// CDCLogFragWriter will so create new fragments if the current fragment is big
// or old enough.
pub struct CDCLogFragWriter {
    config: FragWriterConfig,
    shard_meta: ShardWriterMeta,
    index_meta: IndexMeta,
    frag_path_gen: Arc<FragPathGenerator>,
    cdc_log_receiver: mpsc::Receiver<CDCLogEntry>,
    current_frag_writer: Option<FragWriter>,
    frag_writer_sender: mpsc::Sender<FragUploadTask>,
    shard_tracing_id: String,
}

impl CDCLogFragWriter {
    pub async fn new(
        config: &FragWriterConfig,
        shard_meta: ShardWriterMeta,
        index_meta: IndexMeta,
        frag_path_gen: Arc<FragPathGenerator>,
        cdc_log_receiver: mpsc::Receiver<CDCLogEntry>,
        frag_writer_sender: mpsc::Sender<FragUploadTask>,
    ) -> Result<Self> {
        let shard_tracing_id = shard_meta.shard_tracing_id();

        let frag_path = frag_path_gen.new_path(&shard_meta.index_key).await?;

        let frag_writer = FragWriterBuilder::new()
            .frag_base_path(frag_path_gen.frag_base_path())
            .frag_path(&frag_path)
            .shard_tracing_id(&shard_tracing_id)
            .schema(index_meta.get_tantivy_schema().clone())
            .index_num_threads(config.index_num_threads)
            .index_mem_budget(config.index_mem_budget.0 as usize)
            .is_base(false)
            .analyzer_manager(index_meta.get_analyzer_manager().clone())
            .sort_fields(index_meta.get_sort_fields().to_vec())
            .build()?;

        Ok(CDCLogFragWriter {
            config: config.clone(),
            shard_meta,
            index_meta,
            frag_path_gen,
            cdc_log_receiver,
            current_frag_writer: Some(frag_writer),
            frag_writer_sender,
            shard_tracing_id,
        })
    }

    pub fn take_frag_writer(&mut self) -> Result<FragWriter> {
        self.current_frag_writer
            .take()
            .ok_or_else(|| anyhow!("{} current_frag_writer is none", self.shard_tracing_id))
    }

    // Switch the current fragment writer to a new one.
    // The old fragment writer will be flushed.
    // Return the old fragment writer.
    pub async fn switch_frag_writer(&mut self) -> Result<FragWriter> {
        // Build the new writer first so we never lose the current one if creation
        // fails.
        let new_frag_writer = FragWriterBuilder::new()
            .frag_base_path(self.frag_path_gen.frag_base_path())
            .frag_path(
                &self
                    .frag_path_gen
                    .new_path(&self.shard_meta.index_key)
                    .await?,
            )
            .shard_tracing_id(&self.shard_tracing_id)
            .schema(self.index_meta.get_tantivy_schema().clone())
            .index_num_threads(self.config.index_num_threads)
            .index_mem_budget(self.config.index_mem_budget.0 as usize)
            .is_base(false)
            .analyzer_manager(self.index_meta.get_analyzer_manager().clone())
            .sort_fields(self.index_meta.get_sort_fields().to_vec())
            .build()?;
        let old_frag_writer = self
            .current_frag_writer
            .replace(new_frag_writer)
            .ok_or_else(|| anyhow!("{} current_frag_writer is none", self.shard_tracing_id))?;
        Ok(old_frag_writer)
    }

    pub fn add_tantivy_doc(&mut self, doc: TantivyDoc) -> Result<()> {
        if let Some(frag_writer) = self.current_frag_writer.as_mut() {
            frag_writer.add_tantivy_doc(doc);
            Ok(())
        } else {
            bail!("{} current_frag_writer is none", self.shard_tracing_id);
        }
    }

    async fn set_last_finished_filename(&mut self, filename: CDCFileName) -> Result<()> {
        if let Some(frag_writer) = self.current_frag_writer.as_mut() {
            frag_writer.set_last_cdc_file(Some(filename));
        }
        Ok(())
    }

    fn need_switch_frag_writer(&self) -> bool {
        let frag_writer = self.current_frag_writer.as_ref().unwrap();
        let stat = frag_writer.get_buffered_stat();
        if stat.is_none() {
            return false;
        }
        let stat = stat.unwrap();
        assert!(stat.doc_count > 0, "No docs in frag writer stats");

        let size_trigger = stat.doc_bytes >= self.config.index_flush_size_limit.0 as usize;
        let elapsed = stat.created_time.elapsed();
        let time_trigger = elapsed >= self.config.index_flush_interval.0;

        if size_trigger || time_trigger {
            let elapsed_ms = elapsed.as_millis();
            let label = if size_trigger {
                "size_limit"
            } else {
                "timeout"
            };
            WRITER_FRAG_OPERATION_LATENCY
                .with_label_values(&[format!("write_{}", label).as_str()])
                .observe(elapsed.as_secs_f64());
            info!(
                "switch_frag_writer_triggered, {}, trigger={}, frag_path={}, cdc_last_file={:?}, docs={}, bytes_mib={:.2}, flush_size_limit_mib={:.2}, elapsed_ms={}, flush_interval_ms={}",
                self.shard_tracing_id,
                label,
                frag_writer.frag_path(),
                frag_writer.last_cdc_file(),
                stat.doc_count,
                stat.doc_bytes as f64 / 1024.0 / 1024.0,
                self.config.index_flush_size_limit.as_mb_f64(),
                elapsed_ms,
                self.config.index_flush_interval.as_millis(),
            );
        }
        size_trigger || time_trigger
    }

    // switch frag writer if needed.
    // Sends the old fragment writer to the frag_writer_sender so that it will be
    // uploaded later.
    async fn try_switch_frag_writer(&mut self) -> Result<()> {
        if self.need_switch_frag_writer() {
            let frag_writer = self.switch_frag_writer().await?;
            self.frag_writer_sender
                .send(FragUploadTask::Frag(frag_writer))
                .await?;
        }
        Ok(())
    }

    fn progress_flush_deadline(&self) -> Result<Option<Instant>> {
        let frag_writer = self
            .current_frag_writer
            .as_ref()
            .ok_or_else(|| anyhow!("current_frag_writer is none"))?;
        Ok(frag_writer.progress_flush_deadline(self.config.index_flush_interval.0))
    }

    fn next_flush_deadline(&self) -> Result<Option<Instant>> {
        let frag_buffer_created_time = self
            .current_frag_writer
            .as_ref()
            .ok_or_else(|| anyhow!("current_frag_writer is none"))?
            .get_buffered_stat()
            .map(|s| s.created_time);
        if let Some(created_time) = frag_buffer_created_time {
            let elapsed = created_time.elapsed();
            let max_age = self.config.index_flush_interval.0;
            if elapsed >= max_age {
                Ok(Some(Instant::now()))
            } else {
                Ok(Some(Instant::now() + (max_age - elapsed)))
            }
        } else {
            self.progress_flush_deadline()
        }
    }

    async fn try_submit_progress(&mut self) -> Result<bool> {
        let submit_info = {
            let frag_writer = self
                .current_frag_writer
                .as_mut()
                .ok_or_else(|| anyhow!("current_frag_writer is none"))?;
            let Some(last_cdc_file) = frag_writer.take_progress_file() else {
                return Ok(false);
            };
            SubmitInfo::for_update_progress(last_cdc_file, frag_writer.frag_path())
        };
        self.frag_writer_sender
            .send(FragUploadTask::Progress(submit_info))
            .await?;
        Ok(true)
    }

    async fn handle_flush_deadline(&mut self) -> Result<()> {
        if self.need_switch_frag_writer() {
            self.try_switch_frag_writer().await
        } else {
            let _ = self.try_submit_progress().await?;
            Ok(())
        }
    }

    pub async fn run(&mut self) -> Result<u64> {
        let mut msg_count = 0;

        // Continously listen the cdc_log_receiver until it is closed.
        // If receive a CDCLogEntry, add it to the current fragment writer.
        // If receive a filename, set it as the current fragment writer's commit_seq.
        loop {
            let flush_deadline = self.next_flush_deadline()?;

            tokio::select! {
                biased;  // ​Prioritize handling events from cdc_log_receiver
                msg = self.cdc_log_receiver.recv() => {
                    match msg {
                        Some(cdc_log_entry) => {
                            match cdc_log_entry {
                                CDCLogEntry::CDCData(doc) => {
                                    msg_count += 1;
                                    self.add_tantivy_doc(doc)?;
                                    self.try_switch_frag_writer().await?;
                                }
                                CDCLogEntry::Filename(filename) => {
                                    self.set_last_finished_filename(filename).await?;
                                }
                            }
                        }
                        None => {
                            break;
                        }
                    }
                }
                // If flush_deadline is some, sleep until the deadline and flush the buffer.
                // If flush_deadline is none, this branch will not be executed.
                _ = async {
                    if let Some(deadline) = flush_deadline {
                        tokio::time::sleep_until(deadline).await
                    }
                }, if flush_deadline.is_some() => {
                    self.handle_flush_deadline().await?;
                }
            }
        }
        let frag_writer = self.take_frag_writer()?;
        if let Some(stat) = frag_writer.get_buffered_stat() {
            assert!(stat.doc_count > 0, "No docs in frag writer stats");
            self.frag_writer_sender
                .send(FragUploadTask::Frag(frag_writer))
                .await?;
        } else if let Some(last_cdc_file) = frag_writer.last_cdc_file().clone() {
            self.frag_writer_sender
                .send(FragUploadTask::Progress(SubmitInfo::for_update_progress(
                    last_cdc_file,
                    frag_writer.frag_path(),
                )))
                .await?;
        }
        Ok(msg_count)
    }

    #[cfg(test)]
    pub fn add_raw_doc(&mut self, raw_doc: crate::utils::tantivy_doc::RawDoc) -> Result<()> {
        if self.current_frag_writer.is_none() {
            return Err(anyhow!("current_frag_writer is none"));
        }

        let doc = TantivyDoc::new(&raw_doc, &self.index_meta)?;

        self.current_frag_writer
            .as_mut()
            .unwrap()
            .add_tantivy_doc(doc);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::ops::Bound;

    use rand::Rng;
    use tantivy::{
        Index, Searcher, TantivyDocument, Term,
        collector::DocSetCollector,
        query::RangeQuery,
        schema::{Field, IndexRecordOption, Value},
    };
    use tempfile::TempDir;
    use tici_common::cdc_file_name::CDCFileName;
    use tici_system::readable::ReadableDuration;
    use tokio::sync::mpsc::error::TryRecvError;

    use super::*;
    use crate::{
        data_types::codec::encode_decimal_val,
        test_utils::{
            TEST_DECIMAL_INDEX_META, TEST_FLOAT_INDEX_META, TEST_INDEX_META_OVERLAP,
            TEST_SHARD_META, is_sorted_frag, test_new_frag_path_generator,
        },
        utils::{
            self,
            key::ShardingKey,
            schema::{ColumnInfoExt, HANDLE_STORE_NAME, VERSION_STORE_NAME},
            tantivy_doc::RawDoc,
        },
        writer::cdc_log::{CDCLog, OperationType},
    };

    fn mock_last_cdc_file() -> CDCFileName {
        CDCFileName::parse(
            "tici_default_prefix/cdc/123/461869219056451591/CDC00000000000000000001.json",
        )
        .unwrap()
    }

    #[tokio::test]
    async fn test_new_frag_writer() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        drop(cdc_log_sender);
        let (frag_writer_sender, _) = mpsc::channel(1);
        let cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let schema = cdc_log_frag_writer.index_meta.get_tantivy_schema();

        assert!(schema.get_field(HANDLE_STORE_NAME).is_ok());
        assert!(schema.get_field(VERSION_STORE_NAME).is_ok());
        for col in TEST_INDEX_META_OVERLAP.index_columns() {
            assert!(schema.get_field(&col.store_name()).is_ok());
        }
        for col in TEST_INDEX_META_OVERLAP.pk_columns() {
            assert!(schema.get_field(&col.store_name()).is_ok());
        }
        assert!(cdc_log_frag_writer.current_frag_writer.is_some());
    }

    #[tokio::test]
    async fn test_switch_and_frag_writer() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        drop(cdc_log_sender);
        let (frag_writer_sender, _) = mpsc::channel(1);

        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        cdc_log_frag_writer.switch_frag_writer().await.unwrap();
        cdc_log_frag_writer.switch_frag_writer().await.unwrap();

        cdc_log_frag_writer.take_frag_writer().unwrap();

        assert!(cdc_log_frag_writer.take_frag_writer().is_err());
    }

    #[tokio::test]
    async fn test_add_doc_and_commit() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        // Add a document
        let raw_doc1 = RawDoc {
            handle: ShardingKey::from_raw(b"doc1".to_vec()),
            version: 1,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document One".to_string())),
                (2, Some("This is the first document.".to_string())),
                (3, Some("pk_value1".to_string())),
                (4, Some("pk_value2".to_string())),
            ],
            is_local_timezone: true,
        };
        let raw_doc2 = RawDoc {
            handle: ShardingKey::from_raw(b"doc2".to_vec()),
            version: 2,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document Two".to_string())),
                (2, Some("This is the second document.".to_string())),
                (3, Some("pk_value3".to_string())),
                (4, Some("pk_value4".to_string())),
            ],
            is_local_timezone: true,
        };

        cdc_log_frag_writer.add_raw_doc(raw_doc1).unwrap();
        cdc_log_frag_writer.add_raw_doc(raw_doc2).unwrap();

        // Commit changes
        let frag_path = {
            let mut frag_writer = cdc_log_frag_writer.take_frag_writer().unwrap();
            frag_writer.finalize().unwrap();
            frag_writer.frag_path()
        };

        // Verify that the document was added correctly
        let index = Index::open_in_dir(utils::frag_index_path(&frag_path)).unwrap();
        let reader = index.reader().unwrap();
        let searcher: Searcher = reader.searcher();
        assert_eq!(searcher.num_docs(), 2);

        // Search on fulltext field (ft_2_default) instead of column_2
        let ft_field = cdc_log_frag_writer
            .index_meta
            .get_tantivy_schema()
            .get_field("ft_2_default")
            .unwrap();
        let term = Term::from_field_text(ft_field, "first");
        let query = tantivy::query::TermQuery::new(term, IndexRecordOption::Basic);
        let top_docs = searcher.search(&query, &DocSetCollector).unwrap();
        assert_eq!(top_docs.len(), 1);
    }

    #[tokio::test]
    async fn test_delete_and_update_doc() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        // Insert documents
        let raw_doc1 = RawDoc {
            handle: ShardingKey::from_raw(b"doc1".to_vec()),
            version: 1,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document One".to_string())),
                (2, Some("This is the first document.".to_string())),
                (3, Some("pk_value1".to_string())),
                (4, Some("pk_value2".to_string())),
            ],
            is_local_timezone: true,
        };
        let raw_doc2 = RawDoc {
            handle: ShardingKey::from_raw(b"doc2".to_vec()),
            version: 2,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document Two".to_string())),
                (2, Some("This is the second document.".to_string())),
                (3, Some("pk_value3".to_string())),
                (4, Some("pk_value4".to_string())),
            ],
            is_local_timezone: true,
        };
        cdc_log_frag_writer.add_raw_doc(raw_doc1).unwrap();
        cdc_log_frag_writer.add_raw_doc(raw_doc2).unwrap();

        // Update document
        let raw_doc3 = RawDoc {
            handle: ShardingKey::from_raw(b"doc1".to_vec()),
            version: 3,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document One".to_string())),
                (
                    2,
                    Some("This is the first document that has been updated.".to_string()),
                ),
                (3, Some("pk_value1".to_string())),
                (4, Some("pk_value2".to_string())),
            ],
            is_local_timezone: true,
        };
        cdc_log_frag_writer.add_raw_doc(raw_doc3).unwrap();

        // Delete document
        let raw_doc4 = RawDoc {
            handle: ShardingKey::from_raw(b"doc2".to_vec()),
            version: 4,
            op_type: OperationType::Delete,
            cols: vec![],
            is_local_timezone: true,
        };
        cdc_log_frag_writer.add_raw_doc(raw_doc4).unwrap();

        // Commit changes
        let frag_path = {
            let mut frag_writer = cdc_log_frag_writer.take_frag_writer().unwrap();
            frag_writer.finalize().unwrap();
            frag_writer.frag_path()
        };

        // Verify documents
        let index = Index::open_in_dir(utils::frag_index_path(&frag_path)).unwrap();
        let reader = index.reader().unwrap();
        let searcher: Searcher = reader.searcher();
        assert_eq!(searcher.num_docs(), 1);

        let schema = cdc_log_frag_writer.index_meta.get_tantivy_schema();
        // Search "first" on fulltext field (ft_2_default)
        {
            let ft_field = schema.get_field("ft_2_default").unwrap();
            let term = Term::from_field_text(ft_field, "first");
            let query = tantivy::query::TermQuery::new(term, IndexRecordOption::Basic);
            let docs = searcher.search(&query, &DocSetCollector).unwrap();
            assert_eq!(docs.len(), 1);
            let doc_address = docs.into_iter().next().unwrap();
            let doc: TantivyDocument = searcher.doc(doc_address).unwrap();
            let version = {
                let field = schema.get_field(VERSION_STORE_NAME).unwrap();
                let value = doc.get_first(field).unwrap();
                value.as_u64().unwrap()
            };
            assert_eq!(version, 3);
        }

        // Search "second" on fulltext field (ft_2_default)
        {
            let ft_field = schema.get_field("ft_2_default").unwrap();
            let term = Term::from_field_text(ft_field, "second");
            let query = tantivy::query::TermQuery::new(term, IndexRecordOption::Basic);
            let top_docs = searcher.search(&query, &DocSetCollector).unwrap();
            assert_eq!(top_docs.len(), 0);
        }
    }

    #[tokio::test]
    async fn test_delete_only_and_switch_frag_writer() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        drop(cdc_log_sender);
        let (frag_writer_tx, mut frag_writer_rx) = mpsc::channel(1);
        let frag_writer_config = FragWriterConfig {
            index_flush_interval: ReadableDuration::secs(0),
            ..Default::default()
        };

        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &frag_writer_config,
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_tx,
        )
        .await
        .unwrap();

        // No frag writer switched
        cdc_log_frag_writer.try_switch_frag_writer().await.unwrap();
        assert_eq!(
            TryRecvError::Empty,
            frag_writer_rx.try_recv().err().unwrap()
        );

        let raw_doc = RawDoc {
            handle: ShardingKey::from_raw(b"doc2".to_vec()),
            version: 4,
            op_type: OperationType::Delete,
            cols: vec![],
            is_local_timezone: true,
        };
        cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();

        // One frag writer switched
        cdc_log_frag_writer.try_switch_frag_writer().await.unwrap();
        let mut frag_writer = match frag_writer_rx.try_recv().unwrap() {
            FragUploadTask::Frag(frag_writer) => frag_writer,
            _ => panic!("Unexpected frag upload task"),
        };

        let frag_res = frag_writer.finalize().unwrap();
        assert!(frag_res.seg_ids.is_empty());
    }

    #[tokio::test]
    async fn test_empty_writer_timeout_submits_progress() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        drop(cdc_log_sender);
        let (frag_writer_tx, mut frag_writer_rx) = mpsc::channel(1);
        let frag_writer_config = FragWriterConfig {
            index_flush_interval: ReadableDuration::secs(0),
            ..Default::default()
        };

        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &frag_writer_config,
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_tx,
        )
        .await
        .unwrap();

        let frag_path = cdc_log_frag_writer
            .current_frag_writer
            .as_ref()
            .unwrap()
            .frag_path();
        cdc_log_frag_writer
            .set_last_finished_filename(mock_last_cdc_file())
            .await
            .unwrap();

        cdc_log_frag_writer.handle_flush_deadline().await.unwrap();

        match frag_writer_rx.try_recv().unwrap() {
            FragUploadTask::Progress(_) => {}
            _ => panic!("Unexpected frag upload task"),
        }
        assert_eq!(
            cdc_log_frag_writer
                .current_frag_writer
                .as_ref()
                .unwrap()
                .frag_path(),
            frag_path
        );
        assert!(
            cdc_log_frag_writer
                .current_frag_writer
                .as_ref()
                .unwrap()
                .last_cdc_file()
                .is_none()
        );
        assert!(
            cdc_log_frag_writer
                .current_frag_writer
                .as_ref()
                .unwrap()
                .progress_flush_deadline(frag_writer_config.index_flush_interval.0)
                .is_none()
        );
        assert_eq!(
            frag_writer_rx.try_recv().err().unwrap(),
            TryRecvError::Empty
        );
    }

    #[tokio::test]
    async fn test_pending_progress_is_committed_with_frag_after_doc() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        drop(cdc_log_sender);
        let (frag_writer_tx, mut frag_writer_rx) = mpsc::channel(1);
        let frag_writer_config = FragWriterConfig {
            index_flush_interval: ReadableDuration::secs(0),
            ..Default::default()
        };

        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &frag_writer_config,
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_tx,
        )
        .await
        .unwrap();

        cdc_log_frag_writer
            .set_last_finished_filename(mock_last_cdc_file())
            .await
            .unwrap();
        let raw_doc = RawDoc {
            handle: ShardingKey::from_raw(b"doc1".to_vec()),
            version: 1,
            op_type: OperationType::Insert,
            cols: vec![
                (1, Some("Document One".to_string())),
                (2, Some("This is the first document.".to_string())),
                (3, Some("pk_value1".to_string())),
                (4, Some("pk_value2".to_string())),
            ],
            is_local_timezone: true,
        };
        cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();

        cdc_log_frag_writer.handle_flush_deadline().await.unwrap();

        let frag_writer = match frag_writer_rx.try_recv().unwrap() {
            FragUploadTask::Frag(frag_writer) => frag_writer,
            _ => panic!("Unexpected frag upload task"),
        };
        assert_eq!(frag_writer.last_cdc_file(), &Some(mock_last_cdc_file()));
        assert!(
            cdc_log_frag_writer
                .current_frag_writer
                .as_ref()
                .unwrap()
                .progress_flush_deadline(frag_writer_config.index_flush_interval.0)
                .is_none()
        );
        assert_eq!(
            frag_writer_rx.try_recv().err().unwrap(),
            TryRecvError::Empty
        );
    }

    #[tokio::test]
    async fn test_run_flushes_pending_progress_on_exit() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        drop(cdc_log_sender);
        let (frag_writer_tx, mut frag_writer_rx) = mpsc::channel(1);

        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            TEST_INDEX_META_OVERLAP.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_tx,
        )
        .await
        .unwrap();

        cdc_log_frag_writer
            .set_last_finished_filename(mock_last_cdc_file())
            .await
            .unwrap();
        cdc_log_frag_writer.run().await.unwrap();

        match frag_writer_rx.try_recv().unwrap() {
            FragUploadTask::Progress(_) => {}
            _ => panic!("Unexpected frag upload task"),
        }
    }

    #[tokio::test]
    async fn test_null_column() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);
        let shard_meta = TEST_SHARD_META.clone();
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            shard_meta.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let json = r#"
        {
            "id": 0,
            "database": "test",
            "table": "tp_int",
            "isDdl": false,
            "type": "INSERT",
            "data": [
                {
                    "name": "Document with Nulls",
                    "description": null,
                    "pk1": "pk_value1",
                    "pk2": "pk_value2"
                }
            ],
            "old": null,
            "_tidb": {
                "commitTs": 429918007904436226,
                "rowkey": "dIAAAAAAAAB7X3KAAAAAAAAD7w=="
            }
        }"#;

        let log = CDCLog::new(json)
            .unwrap()
            .expect("Expected DML event in test");
        log.get_sharding_key_pair(&index_meta).unwrap();
        let doc = TantivyDoc::new(&log, &index_meta).unwrap();
        cdc_log_frag_writer.add_tantivy_doc(doc).unwrap();

        // Commit changes
        let frag_path = {
            let mut frag_writer = cdc_log_frag_writer.take_frag_writer().unwrap();
            frag_writer.finalize().unwrap();
            frag_writer.frag_path()
        };

        // Verify that the document was added correctly
        let index = Index::open_in_dir(utils::frag_index_path(&frag_path)).unwrap();
        let reader = index.reader().unwrap();
        let searcher: Searcher = reader.searcher();
        assert_eq!(searcher.num_docs(), 1);

        let schema = cdc_log_frag_writer.index_meta.get_tantivy_schema();
        // Search on fulltext field (ft_1_default)
        let ft_field = schema.get_field("ft_1_default").unwrap();
        let term = Term::from_field_text(ft_field, "nulls");
        let query = tantivy::query::TermQuery::new(term, IndexRecordOption::Basic);
        let docs = searcher.search(&query, &DocSetCollector).unwrap();
        assert_eq!(docs.len(), 1);
        let desc_field = schema.get_field("column_2").unwrap();
        for addr in docs {
            let doc: TantivyDocument = searcher.doc(addr).unwrap();
            assert!(doc.get_first(desc_field).is_none());
        }
    }

    #[tokio::test]
    async fn test_table_id_mismatch() {
        let index_meta = TEST_INDEX_META_OVERLAP.clone();

        // Table id from the rowkey was 118, mismatch with the index_meta table_id 123
        let json = r#"
        {
            "id": 0,
            "database": "test",
            "table": "tp_int",
            "isDdl": false,
            "type": "INSERT",
            "data": [
                {
                    "name": "Document with Nulls",
                    "description": null,
                    "pk1": "pk_value1",
                    "pk2": "pk_value2"
                }
            ],
            "old": null,
            "_tidb": {
                "commitTs": 429918007904436226,
                "rowkey": "dIAAAAAAAAB2X3KAAAAAAAAD7w=="
            }
        }"#;

        let log = CDCLog::new(json)
            .unwrap()
            .expect("Expected DML event in test");
        assert!(log.get_sharding_key_pair(&index_meta).is_err());
    }

    #[tokio::test]
    async fn test_float() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);
        let index_meta = TEST_FLOAT_INDEX_META.clone();
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let mut rng = rand::rng();
        let mut rands = vec![];
        for i in 0..100 {
            let score1: f32 = rng.random_range(1.0..100.0);
            let score2: f64 = rng.random_range(100.0..200.0);
            rands.push((i, score1, score2));
        }

        for (ver, (id, score1, score2)) in rands.iter().enumerate() {
            let raw_doc = RawDoc {
                handle: ShardingKey::from_raw(id.to_string().into_bytes()),
                version: ver as u64 + 1,
                op_type: OperationType::Insert,
                cols: vec![
                    (1, Some(id.to_string())),
                    (2, Some(score1.to_string())),
                    (3, Some(score2.to_string())),
                ],
                is_local_timezone: true,
            };
            cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
        }

        // Commit changes
        let frag_path = {
            let mut frag_writer = cdc_log_frag_writer.take_frag_writer().unwrap();
            frag_writer.finalize().unwrap();
            frag_writer.frag_path()
        };

        assert!(is_sorted_frag(&frag_path, &index_meta));

        // Verify that the document was added correctly
        let index = Index::open_in_dir(utils::frag_index_path(&frag_path)).unwrap();
        let reader = index.reader().unwrap();
        let searcher: Searcher = reader.searcher();
        assert_eq!(searcher.num_docs(), 100);

        let schema = cdc_log_frag_writer.index_meta.get_tantivy_schema();
        let id_field = schema.get_field("column_1").unwrap();
        let score1_field = schema.get_field("column_2").unwrap();
        let score2_field = schema.get_field("column_3").unwrap();

        let mut rand_range = || {
            let s1: f64 = rng.random_range(1.0..1000.0);
            let s2: f64 = rng.random_range(1.0..1000.0);
            let start = s1.min(s2);
            let end = s1.max(s2);
            (start, end)
        };

        let float_range_query = |field: Field, start: f64, end: f64| {
            let lower_term = Term::from_field_f64(field, start);
            let upper_term = Term::from_field_f64(field, end);
            let range_query =
                RangeQuery::new(Bound::Included(lower_term), Bound::Included(upper_term));
            let doc_addrs = searcher.search(&range_query, &DocSetCollector).unwrap();
            let mut ids = vec![];
            for doc_addr in doc_addrs {
                let doc: TantivyDocument = searcher.doc(doc_addr).unwrap();
                let id: i64 = doc.get_first(id_field).unwrap().as_i64().unwrap();
                ids.push(id);
            }
            ids.sort();
            ids
        };

        for _ in 0..10 {
            let (start, end) = rand_range();

            let mut score1_ids = vec![];
            for (id, score1, _) in rands.iter() {
                if *score1 as f64 >= start && *score1 as f64 <= end {
                    score1_ids.push(*id as i64);
                }
            }
            score1_ids.sort();
            let ids1 = float_range_query(score1_field, start, end);
            assert_eq!(score1_ids, ids1);

            let mut score2_ids = vec![];
            for (id, _, score2) in rands.iter() {
                if *score2 >= start && *score2 <= end {
                    score2_ids.push(*id as i64);
                }
            }
            score2_ids.sort();
            let ids2 = float_range_query(score2_field, start, end);
            assert_eq!(score2_ids, ids2);
        }
    }

    #[tokio::test]
    async fn test_decimal() {
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1);
        let (frag_writer_sender, _) = mpsc::channel(1);
        drop(cdc_log_sender);
        let index_meta = TEST_DECIMAL_INDEX_META.clone();
        let mut cdc_log_frag_writer = CDCLogFragWriter::new(
            &FragWriterConfig::default(),
            TEST_SHARD_META.clone(),
            index_meta.clone(),
            test_new_frag_path_generator(local_data_path),
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await
        .unwrap();

        let mut rng = rand::rng();
        let mut rands = vec![];
        for i in 0..100 {
            let score1: f64 = rng.random_range(-100.0..100.0);
            let score2: f64 = rng.random_range(-200.0..200.0);
            rands.push((i, score1, score2));
        }

        // Deterministic edge-case docs: near zero and a unique value for single-hit
        // test
        let edge_docs: [(i32, f64, f64); 4] = [
            (100, -0.01, 0.0),    // score1 negative near zero
            (101, 0.0, 0.01),     // score1 zero, score2 positive near zero
            (102, 0.01, -0.01),   // score1 positive near zero, score2 negative near zero
            (200, 12345.67, 0.0), // unique value for single-doc range test
        ];
        for (id, score1, score2) in edge_docs {
            rands.push((id, score1, score2));
        }

        for (ver, (id, score1, score2)) in rands.iter().enumerate() {
            let raw_doc = RawDoc {
                handle: ShardingKey::from_raw(id.to_string().into_bytes()),
                version: ver as u64 + 1,
                op_type: OperationType::Insert,
                cols: vec![
                    (1, Some(id.to_string())),
                    (2, Some(score1.to_string())),
                    (3, Some(score2.to_string())),
                ],
                is_local_timezone: true,
            };
            cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap();
        }

        // Commit changes
        let frag_path = {
            let mut frag_writer = cdc_log_frag_writer.take_frag_writer().unwrap();
            frag_writer.finalize().unwrap();
            frag_writer.frag_path()
        };

        assert!(is_sorted_frag(&frag_path, &index_meta));

        // Verify that the document was added correctly
        let index = Index::open_in_dir(utils::frag_index_path(&frag_path)).unwrap();
        let reader = index.reader().unwrap();
        let searcher: Searcher = reader.searcher();
        assert_eq!(searcher.num_docs(), 104, "100 random + 4 edge-case docs");

        let schema = index_meta.get_tantivy_schema();
        let id_field = schema.get_field("column_1").unwrap();
        let score1_field = schema.get_field("column_2").unwrap();
        let score2_field = schema.get_field("column_3").unwrap();

        let decimal_range_query = |field: Field, start: &[u8], end: &[u8]| {
            let lower_term = Term::from_field_bytes(field, start);
            let upper_term = Term::from_field_bytes(field, end);
            let range_query =
                RangeQuery::new(Bound::Included(lower_term), Bound::Included(upper_term));
            let doc_addrs = searcher.search(&range_query, &DocSetCollector).unwrap();
            let mut ids = vec![];
            for doc_addr in doc_addrs {
                let doc: TantivyDocument = searcher.doc(doc_addr).unwrap();
                let id: i64 = doc.get_first(id_field).unwrap().as_i64().unwrap();
                ids.push(id);
            }
            ids.sort();
            ids
        };

        let column_infos = index_meta.get_column_infos();
        let encode_range = |col_idx: usize, start: f64, end: f64| -> (Vec<u8>, Vec<u8>) {
            let mut start_encoded = vec![];
            encode_decimal_val(
                &mut start_encoded,
                &start.to_string(),
                column_infos[col_idx].column_length as u8,
                column_infos[col_idx].decimal as u8,
            )
            .unwrap();
            let mut end_encoded = vec![];
            encode_decimal_val(
                &mut end_encoded,
                &end.to_string(),
                column_infos[col_idx].column_length as u8,
                column_infos[col_idx].decimal as u8,
            )
            .unwrap();
            (start_encoded, end_encoded)
        };

        // Deterministic: empty result
        let (start_enc, end_enc) = encode_range(1, 50000.0, 60000.0);
        let empty_ids = decimal_range_query(score1_field, &start_enc, &end_enc);
        assert!(
            empty_ids.is_empty(),
            "range [50000, 60000] should match no doc"
        );

        // Deterministic: single-doc hit (only id 200 has score1 = 12345.67)
        let (start_enc, end_enc) = encode_range(1, 12345.0, 12346.0);
        let single_ids = decimal_range_query(score1_field, &start_enc, &end_enc);
        assert_eq!(
            single_ids,
            [200],
            "range [12345, 12346] should match only id 200"
        );

        // Deterministic: negative-only range for score1
        let mut expected_neg = vec![];
        for (id, score1, _) in rands.iter() {
            if *score1 >= -10.0 && *score1 <= -1.0 {
                expected_neg.push(*id as i64);
            }
        }
        expected_neg.sort();
        let (start_enc, end_enc) = encode_range(1, -10.0, -1.0);
        let neg_ids = decimal_range_query(score1_field, &start_enc, &end_enc);
        assert_eq!(neg_ids, expected_neg, "score1 in [-10, -1]");

        // Deterministic: cross-zero range for score1
        let mut expected_cross = vec![];
        for (id, score1, _) in rands.iter() {
            if *score1 >= -1.0 && *score1 <= 1.0 {
                expected_cross.push(*id as i64);
            }
        }
        expected_cross.sort();
        let (start_enc, end_enc) = encode_range(1, -1.0, 1.0);
        let cross_ids = decimal_range_query(score1_field, &start_enc, &end_enc);
        assert_eq!(cross_ids, expected_cross, "score1 in [-1, 1] (cross zero)");

        // Random range queries for score1 and score2
        let mut rand_range = || {
            let s1: f64 = rng.random_range(-1000.0..1000.0);
            let s2: f64 = rng.random_range(-1000.0..1000.0);
            let start = s1.min(s2);
            let end = s1.max(s2);
            (start, end)
        };

        for _ in 0..10 {
            let (start, end) = rand_range();

            let mut score1_ids = vec![];
            for (id, score1, _) in rands.iter() {
                if *score1 >= start && *score1 <= end {
                    score1_ids.push(*id as i64);
                }
            }
            score1_ids.sort();

            let (start_encoded, end_encoded) = encode_range(1, start, end);
            let ids1 = decimal_range_query(score1_field, &start_encoded, &end_encoded);
            assert_eq!(score1_ids, ids1, "start={}, end={}", start, end);
        }

        for _ in 0..10 {
            let (start, end) = rand_range();

            let mut score2_ids = vec![];
            for (id, _, score2) in rands.iter() {
                if *score2 >= start && *score2 <= end {
                    score2_ids.push(*id as i64);
                }
            }
            score2_ids.sort();

            let (start_encoded, end_encoded) = encode_range(2, start, end);
            let ids2 = decimal_range_query(score2_field, &start_encoded, &end_encoded);
            assert_eq!(score2_ids, ids2);
        }
    }
}