use std::collections::HashSet;

use anyhow::Result;
use log::info;
use tici::config::{make_default_test_meta_config, make_default_test_worker_config};
use tici_common::{index::IndexKey, keyspace::NULLSPACEID};
use tici_test_utils::test_utils::test_suite::E2ETestSuite;

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

    rt.block_on(async {
        let meta_cfg = make_default_test_meta_config(true);
        let writer_cfg = make_default_test_worker_config(true);
        let mut test_suite = E2ETestSuite::new_with_cfg(meta_cfg.clone(), writer_cfg, true)
            .await
            .unwrap();
        test_suite.wait_heartbeat().await;

        let keyspace_id = NULLSPACEID;
        let index_id = 123;
        let table_id = 1001;
        let physical_table_ids = vec![table_id + 1, table_id + 2, table_id + 3, table_id + 4];

        info!(
            "Start partition table integration test, table_id={}, index_id={}, physical_table_ids={:?} ...",
            table_id, index_id, physical_table_ids
        );

        // Create partition table index
        let logical_index_key = IndexKey::new(keyspace_id, table_id, index_id);
        test_suite
            .must_create_partition_index(&logical_index_key, &physical_table_ids)
            .await
            .unwrap();

        // Parent index should not be found
        assert!(
            test_suite
                .get_meta_service()
                .get_index(&logical_index_key)
                .await
                .is_err()
        );

        // Check if the index is created correctly
        let worker_addr = test_suite.get_one_worker_node_addr();
        for physical_table_id in physical_table_ids.iter().copied() {
            let index_key = IndexKey::new(keyspace_id, physical_table_id, index_id);
            let index = test_suite.get_meta_service().get_index(&index_key).await;
            assert!(index.is_ok());
            let shards = test_suite.collect_shards_from_index(&index_key).await;
            assert_eq!(shards.len(), 1);
            assert_eq!(shards[0].writer, worker_addr);
        }

        // Check if the indexes are restored correctly
        let indexes = test_suite.collect_indexes_from_restore().await;
        assert_eq!(indexes.len(), physical_table_ids.len());
        let actual_tables: HashSet<i64> = indexes
            .into_iter()
            .map(|(index, ..)| index.table_id)
            .collect();
        let expected_tables: HashSet<i64> = physical_table_ids.iter().copied().collect();
        assert_eq!(actual_tables, expected_tables);

        // Check if the shards are restored correctly
        let shards = test_suite.collect_shards_from_restore().await;
        assert_eq!(shards.len(), physical_table_ids.len());
        let shard_tables: HashSet<i64> = shards.iter().map(|s| s.table_id()).collect();
        assert_eq!(shard_tables, expected_tables);

        // Check if drop_index works correctly (by parent table id)
        test_suite
            .must_drop_index(&logical_index_key)
            .await
            .unwrap();
        for physical_table_id in physical_table_ids.iter() {
            let index_key = IndexKey::new(keyspace_id, *physical_table_id, index_id);
            assert!(
                test_suite
                    .get_meta_service()
                    .get_index(&index_key)
                    .await
                    .is_err()
            );
        }

        for (_, reader) in test_suite.readers.iter_mut() {
            reader.stop().await;
        }
        test_suite.stop().await.unwrap();
    });

    Ok(())
}