use std::{sync::Arc, time::Duration};

use anyhow::Result;
use log::{error, info, warn};
use tokio::{sync::mpsc, task::JoinHandle};

use crate::{
    ShardWriterMeta,
    config::WorkerNodeConfig,
    metadata::metadata_store::MetadataStore,
    utils::{
        frag_path_generator::FragPathGenerator, meta_client_pool::MetaClientPool, schema::IndexMeta,
    },
    writer::{
        cdc_log::CDCLogEntry, cdc_log_frag_writer::CDCLogFragWriter,
        frag_batch_uploader::FragBatchUploader, frag_uploader::FragUploader,
    },
};

pub struct ShardWriter {
    shard_meta: ShardWriterMeta,
    cdc_log_frag_writer_handle: JoinHandle<Result<u64>>,
    frag_uploader: FragUploader,
}

impl ShardWriter {
    #[allow(clippy::too_many_arguments)]
    pub async fn new(
        shard_meta: &ShardWriterMeta,
        index_meta: &IndexMeta,
        config: &WorkerNodeConfig,
        frag_batch_uploader: Arc<FragBatchUploader>,
        meta_client_pool: Arc<dyn MetaClientPool>,
        frag_path_gen: Arc<FragPathGenerator>,
        cdc_log_receiver: mpsc::Receiver<CDCLogEntry>,
        metadata_store: Arc<MetadataStore>,
    ) -> Result<ShardWriter> {
        let (frag_writer_sender, frag_writer_receiver) = mpsc::channel(100);

        let cdc_log_frag_writer = CDCLogFragWriter::new(
            &config.frag_writer,
            shard_meta.clone(),
            index_meta.clone(),
            frag_path_gen,
            cdc_log_receiver,
            frag_writer_sender,
        )
        .await?;

        let tracing_id = shard_meta.shard_tracing_id();
        let frag_uploader = FragUploader::new(
            tracing_id.clone(),
            &shard_meta.index_key,
            shard_meta.shard_header.shard_id,
            config.server.get_advertise_addr().to_string(),
            meta_client_pool,
            frag_batch_uploader,
            shard_meta.frag_location.bucket.clone(),
            shard_meta.frag_location.prefix.clone(),
            frag_writer_receiver,
            true,
            metadata_store,
        );

        let cdc_log_frag_writer_handle = tokio::spawn(async move {
            let mut cdc_log_frag_writer = cdc_log_frag_writer;
            loop {
                let result = cdc_log_frag_writer.run().await;
                match result {
                    Ok(msg_count) => {
                        info!(
                            "cdc_log_frag_writer_exit, {}, status=normal, msg_count={}",
                            tracing_id, msg_count
                        );
                        return Ok(msg_count);
                    }
                    Err(e) => {
                        error!(
                            "cdc_log_frag_writer_exit, {}, status=error_retry, retry_delay_s=1, err={:?}",
                            tracing_id, e
                        );
                        tokio::time::sleep(Duration::from_secs(1)).await;
                    }
                }
            }
        });

        let shard_writer = ShardWriter {
            shard_meta: shard_meta.clone(),
            cdc_log_frag_writer_handle,
            frag_uploader,
        };

        Ok(shard_writer)
    }

    pub async fn shutdown(self) {
        let tracing_id = self.shard_meta.shard_tracing_id();
        self.cdc_log_frag_writer_handle.abort();
        match self.cdc_log_frag_writer_handle.await {
            Ok(Ok(msg_count)) => info!(
                "cdc_log_frag_writer_shutdown, {}, status=finished_before_abort, msg_count={}",
                tracing_id, msg_count
            ),
            Ok(Err(e)) => warn!(
                "cdc_log_frag_writer_shutdown, {}, status=error_before_abort, err={:?}",
                tracing_id, e
            ),
            Err(e) if e.is_cancelled() => info!(
                "cdc_log_frag_writer_shutdown, {}, status=cancelled",
                tracing_id
            ),
            Err(e) => warn!(
                "cdc_log_frag_writer_shutdown, {}, status=join_failed, err={:?}",
                tracing_id, e
            ),
        }

        self.frag_uploader.shutdown().await;
    }

    pub fn shard_meta(&self) -> &ShardWriterMeta {
        &self.shard_meta
    }
}

#[cfg(test)]
mod tests {

    use std::{collections::BTreeMap, ops::Bound, sync::Arc};

    use rand::Rng;
    use tantivy::{
        DocAddress, Index, Searcher, TantivyDocument, Term,
        collector::DocSetCollector,
        query::{AllQuery, RangeQuery, TermQuery},
        schema::{Field, Value},
    };
    use tempfile::TempDir;
    use tici_common::index::IndexKey;
    use tici_grpc::KeyRange;
    use tici_test_utils::test_utils::init_logger_for_ut;

    use crate::{
        config::WorkerNodeConfig,
        test_utils::*,
        utils::{
            self,
            meta_client_pool::{MetaClientPool, MockMetaClientPool},
            s3_frag,
            schema::ColumnInfoExt,
        },
    };

    #[tokio::test]
    async fn test_add_doc() {
        init_logger_for_ut();
        let shard_meta = test_get_shard_writer_meta();
        let index_meta = TEST_HDFSLOGS_INDEX_META.clone();
        let temp_dir = TempDir::new().unwrap();
        let config = WorkerNodeConfig::new_for_test(temp_dir.path().to_str().unwrap());
        let backend_storage = test_create_backend_storage().await.unwrap();

        let meta_client_pool = Arc::new(MockMetaClientPool::new());
        write_test_data(
            &config,
            &[shard_meta.clone()],
            &index_meta,
            backend_storage.clone(),
            meta_client_pool.clone(),
        )
        .await
        .unwrap();

        let frag_parent_path = format!(
            "{}/{}",
            config.storage.frag_base_path(),
            tici_common::s3_keys::index_key(
                shard_meta.keyspace_id(),
                shard_meta.table_id(),
                shard_meta.index_id()
            )
        );
        // frag parent path is empty
        assert!(
            std::fs::read_dir(frag_parent_path)
                .unwrap()
                .next()
                .is_none()
        );

        let mut meta_client = meta_client_pool.get_meta_client().await.unwrap();
        let (frag_metas, _seq) = meta_client
            .get_frags_by_keyrange(
                shard_meta.shard_header.shard_id,
                &IndexKey::new_nullspace(0, 0),
                KeyRange::default(),
            )
            .await
            .unwrap();
        let temp_dir = TempDir::new().unwrap();
        let local_data_path = temp_dir.path().to_str().unwrap();
        let mut frag_paths = Vec::with_capacity(frag_metas.len());
        for frag_meta in frag_metas {
            let frag_path = s3_frag::download_frag(
                backend_storage.as_ref(),
                &shard_meta.frag_location.bucket,
                &frag_meta,
                local_data_path,
            )
            .await
            .unwrap();
            assert!(is_sorted_frag(&frag_path, &index_meta));
            frag_paths.push(frag_path)
        }

        let schema = index_meta.get_tantivy_schema();
        let column_infos = index_meta.get_column_infos();
        let id_field = schema.get_field(&column_infos[2].store_name()).unwrap();
        let tenant_id_field = schema.get_field(&column_infos[3].store_name()).unwrap();
        let mut num_docs = 0;
        for frag_path in frag_paths {
            let index = Index::open_in_dir(utils::frag_index_path(&frag_path)).unwrap();
            let reader = index.reader().unwrap();
            let searcher: Searcher = reader.searcher();
            num_docs += searcher.num_docs();

            let mut ids = BTreeMap::new();
            let mut tenant_ids = BTreeMap::new();
            let doc_addrs = searcher.search(&AllQuery, &DocSetCollector).unwrap();
            for doc_addr in doc_addrs {
                let retrieved_doc: TantivyDocument = searcher.doc(doc_addr).unwrap();
                let id = retrieved_doc.get_first(id_field).unwrap().as_i64().unwrap();
                ids.entry(id).or_insert_with(Vec::new).push(doc_addr);

                let tenant_id = retrieved_doc
                    .get_first(tenant_id_field)
                    .unwrap()
                    .as_i64()
                    .unwrap();
                tenant_ids
                    .entry(tenant_id)
                    .or_insert_with(Vec::new)
                    .push(doc_addr);
            }

            let verify_integer_index = |ids: BTreeMap<i64, Vec<DocAddress>>, field: Field| {
                for (id, docs) in ids.iter() {
                    let term = Term::from_field_i64(field, *id);
                    let term_query =
                        TermQuery::new(term, tantivy::schema::IndexRecordOption::Basic);
                    let doc_addrs = searcher.search(&term_query, &DocSetCollector).unwrap();
                    assert_eq!(docs.len(), doc_addrs.len());
                    for doc in docs {
                        assert!(doc_addrs.contains(doc));
                    }
                }
                let all_ids: Vec<i64> = ids.keys().copied().collect();
                let mut rng = rand::rng();

                for _ in 0..5 {
                    if all_ids.is_empty() {
                        break;
                    }

                    let idx1 = rng.random_range(0..all_ids.len());
                    let idx2 = rng.random_range(0..all_ids.len());
                    let start = all_ids[idx1].min(all_ids[idx2]);
                    let end = all_ids[idx1].max(all_ids[idx2]);

                    let mut expected_doc_addrs = Vec::new();
                    for id in all_ids.iter() {
                        if *id >= start && *id <= end {
                            expected_doc_addrs.extend(ids.get(id).unwrap());
                        }
                    }

                    let lower_term = Term::from_field_i64(field, start);
                    let upper_term = Term::from_field_i64(field, end);

                    let range_query =
                        RangeQuery::new(Bound::Included(lower_term), Bound::Included(upper_term));
                    let actual_doc_addrs = searcher.search(&range_query, &DocSetCollector).unwrap();
                    assert_eq!(expected_doc_addrs.len(), actual_doc_addrs.len());
                    for doc in &expected_doc_addrs {
                        assert!(actual_doc_addrs.contains(doc));
                    }
                }
            };
            verify_integer_index(ids, id_field);
            verify_integer_index(tenant_ids, tenant_id_field);
        }
        assert_eq!(num_docs, 150000);
    }
}