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

use anyhow::Result;
use log::{info, warn};
use storage_backend::storage::create_backend_storage_for_test;
use tici::config::{make_default_test_meta_config, make_default_test_worker_config};
use tici_common::{
    codec::{codec_v2::remove_keyspace_prefix, table_end_key, table_start_key},
    index::IndexKey,
};
use tici_grpc::{DropIndexRequest, KeyRange, ShardManifestHeader};
use tici_search_lib::{
    count_by_reader,
    ffi::{CountResult, Range, Shard},
};
use tici_shard::reader::shard_mirror_manager::ShardMirrorManager;
use tici_system::readable::{ReadableDuration, ReadableSize};
use tici_test_utils::test_utils::{
    reader_handle::make_default_reader_config, test_suite::E2ETestSuite, *,
};

use crate::e2e::{Checker, Measurement, build_default_query_expr, get_current_mem};

#[derive(PartialEq)]
pub enum RunConfig {
    Simple,
    Workload,
}

async fn count_for_shard(
    reader: Arc<ShardMirrorManager>,
    shard_header: &ShardManifestHeader,
    index_key: &IndexKey,
) -> Result<CountResult> {
    let query = build_default_query_expr();
    let read_ts = 0_u64;
    let shard = Shard {
        keyspace_id: index_key.keyspace_id,
        index_id: index_key.index_id,
        shard_id: shard_header.shard_id,
        shard_epoch: shard_header.epoch,
    };
    let ranges = vec![Range {
        start: remove_keyspace_prefix(shard_header.start_key.as_slice()),
        end: remove_keyspace_prefix(shard_header.end_key.as_slice()),
    }];
    let shard_reader = reader
        .get_shard(shard_header.shard_id)
        .ok_or_else(|| anyhow::anyhow!("Shard {} not found", shard_header.shard_id))?
        .get_reader();
    let res = count_by_reader(shard_reader, &shard, &ranges, &query, &read_ts)
        .map_err(|e| anyhow::anyhow!("{}", e))?;
    Ok(res)
}

#[test]
fn test_main() -> Result<()> {
    let allocate_parallelism = 2;
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4 + allocate_parallelism)
        .enable_all()
        .build()?;

    let (max_att, run_config) = match std::env::var("TEST_SIZE") {
        Ok(value) => {
            let w: u64 = value.parse().expect("TEST_SIZE must be a number");
            (w, RunConfig::Workload)
        }
        Err(_) => (100, RunConfig::Simple),
    };

    info!("work dir {}", std::env::current_dir().unwrap().display());

    let mut meta_cfg = make_default_test_meta_config(true);
    let mut writer_cfg = make_default_test_worker_config(true);

    meta_cfg.cache_scheduling.allocate_parallelism = 2;
    meta_cfg.shard.gc_config.gc_worker_check_interval = ReadableDuration::secs(20000);
    let max_size = if run_config == RunConfig::Simple {
        meta_cfg.shard.max_size = ReadableSize::mb(50);
        meta_cfg.shard.split.split_threshold = 0.2;
        writer_cfg.frag_writer.index_flush_size_limit = ReadableSize::mb(5);
        0
    } else {
        let index_flush_size_limit = 32;
        meta_cfg.shard.max_size = ReadableSize::mb(512);
        writer_cfg.frag_writer.index_flush_size_limit = ReadableSize::mb(index_flush_size_limit);
        writer_cfg.frag_writer.index_mem_budget = ReadableSize::mb(1000);
        writer_cfg.frag_writer.index_num_threads = 1;
        5 * 1024 * 1024 * 1024
    };

    info!("using test cfg={:?}", meta_cfg);

    rt.block_on(async {
        let tidb_server = meta_cfg.tidb_server.dsns.first().unwrap();
        workload_ci::bootstrap_empty_hdfs_table(tidb_server.as_str(), "hdfs_10w")
            .await
            .unwrap();
        let table_id = workload_ci::get_table_id(tidb_server.as_str()).await.unwrap();
        let index_id = 1001;
        info!("Start worker integration test, table_id={} ...", table_id);
        let mut test_suite = E2ETestSuite::new_with_cfg(meta_cfg.clone(), writer_cfg, true).await.unwrap();
        // wait heartbeat for worker registered
        test_suite.wait_heartbeat().await;

        let searchlib_config = make_default_reader_config(&test_suite.get_meta_service().config.load_full(), "127.0.0.1:19410".to_string(), "127.0.0.1:19411".to_string());
        test_suite.add_real_reader_node_with_cfg(searchlib_config).await;

        let index_key = IndexKey::new_nullspace(table_id, index_id);
        info!("Creating index with index_key={:?}", index_key);
        test_suite.must_create_index(&index_key).await.unwrap();
        info!("Inserting HDFS logs after create_index for index_key={:?}", index_key);
        workload_ci::insert_hdfs_logs(tidb_server.as_str(), "hdfs_10w", 100_000)
            .await
            .unwrap();

        info!("Validate add index result for index_key={:?}", index_key);
        test_suite.wait_heartbeat().await;
        test_suite.validate_add_index_result(&index_key).await;

        must_finally(async || {
            let res = test_suite
                .send_get_shard_local_cache_request(
                    &index_key,
                    vec![KeyRange {
                        start_key: table_start_key(&index_key).to_vec(),
                        end_key: table_end_key(&index_key).to_vec(),
                    }],
                )
                .await
                .unwrap();
            test_suite.check_valid_shard_allocation(&res, 0).is_ok()
        })
        .await;

        let storage =
            create_backend_storage_for_test(&test_suite.get_one_worker_node_handle().config.s3)
                .await
                .unwrap();

        let mut att = 0;
        let start = Instant::now();
        let mut measurement = Measurement::default();
        let checker = Checker::new(max_size);
        let mut real = false;

        loop {
            let index = test_suite.get_index(&index_key).await;

            let current_stats = Checker::get_current_stat(&test_suite, &meta_cfg, &index_key).await;
            real |= must_check_gc_files(
                test_suite.get_meta_service().shard_task_manager.as_ref(),
                index.clone(),
                &storage,
            )
            .await
            .unwrap();
            if checker.is_ok(&current_stats, &test_suite) && real
            {
                info!(
                    "finished: current({}/{}), stats={:?}, index_stats=({}), elapsed={:?} s",
                    att,
                    max_att,
                    current_stats,
                    index.stats.to_short_string(),
                    start.elapsed().as_secs(),
                );
                break;
            }
            if att > max_att {
                if run_config == RunConfig::Simple {
                    panic!(
                        "exceed current_stats={:?}, index_stats=({})",
                        current_stats,
                        index.stats.to_short_string(),
                    );
                } else {
                    break;
                }
            } else {
                info!(
                    "pending: current({}/{}), {}, current_stats={:?}, index_stats=({}), elapsed={:?} s",
                    att,
                    max_att,
                    measurement.get_stats_string(&current_stats, index.clone()),
                    current_stats,
                    index.stats.to_short_string(),
                    start.elapsed().as_secs(),
                );
            }
            std::thread::sleep(std::time::Duration::from_millis(3_000));
            att += 1;
        }

        // Everytime we append, it will add `default_rows_batch` rows, which contribute to 1 in `rows_from_frag`.
        warn!("Final result: Mem={}MB, Time={}s, total_size={:.2}MB, total_shards={}",
            get_current_mem().as_mb_f64(),
            start.elapsed().as_secs(),
            test_suite.compute_total_frags(&index_key).await,
            test_suite.collect_shards_from_restore().await.len(),
        );

        info!("Stop worker");
        test_suite.stop_worker().await.unwrap();
        info!("Stop worker finished");

        // Stop the workload, and check eventual consistency.
        info!("Start Validation ...");
        must_finally(async || {
            let res = test_suite
                .send_get_shard_local_cache_request(&index_key, vec![KeyRange {
                    start_key: table_start_key(&index_key).to_vec(),
                    end_key: table_end_key(&index_key).to_vec(),
                }])
                .await
                .unwrap();
            test_suite.check_valid_shard_allocation(&res, 0).unwrap();
            let shards = test_suite.collect_shards_from_restore().await;
            for shard in shards.iter() {
                if !test_suite
                    .check_shard_allocated_on_reader_node(&shard.shard_id, "127.0.0.1:19410")
                    .await {
                        return false;
                    }
                if !test_suite
                    .check_get_shard_manifest(
                        &index_key,
                        shard.shard_id,
                        shard.seq,
                        shard.manifest.get_fragments().iter().map(|x| x.frag_path.clone()).collect::<Vec<_>>(),
                    )
                    .await {
                        return false;
                    }
            }

            let mut total_count = 0_u64;
            for shard_info in res.shard_local_cache_infos.iter() {
                let shard_header = shard_info.shard.as_ref().expect("missing shard header");
                let reader = test_suite.get_reader("127.0.0.1:19410").unwrap_right().read().await.get_service().shard_mirror_manager.clone();
                let count_res = match count_for_shard(reader, shard_header, &index_key).await {
                    Ok(ok) => ok,
                    Err(_) => {
                        // shard not found
                        return false;
                    }
                };
                total_count += count_res.count;
            }
            info!("have found {} records", total_count);
            total_count > 0
        }).await;
        info!("Finished Validation ...");

        meta_cfg.shard.enable_split = false;
        meta_cfg.shard.auto_compaction = false;
        meta_cfg.server.status_addr = String::new();
        info!("Restart meta");
        test_suite.restart_meta(meta_cfg.clone()).await.unwrap();

        let drop_index_req = DropIndexRequest { table_id, index_id, keyspace_id: index_key.keyspace_id };

        let mut client = test_suite.get_meta_client();
        let _ = client
            .drop_index(drop_index_req.clone())
            .await.unwrap();

        assert_eq!(test_suite.collect_indexes_from_restore().await.len(), 0);

        assert_eq!(
            test_suite
                .collect_worker_nodes()
                .await
                .first()
                .unwrap()
                .1
                .shards_len()
                .await,
            0
        );
        test_suite.stop_readers().await.unwrap();
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}