use std::{error::Error, sync::Arc};

use anyhow::Result;
use log::info;
use tici::config::{make_default_test_meta_config, make_default_test_worker_config};
use tici_common::{
    codec::{table_end_key, table_start_key},
    index::IndexKey,
    keyspace::NULLSPACEID,
};
use tici_grpc::{KeyRange, ShardManifestHeader};
use tici_search_lib::ffi::Range;
use tici_shard::reader::shard_mirror_manager::ShardMirrorManager;
use tici_test_utils::test_utils::{
    reader_handle::make_default_reader_config, test_suite::E2ETestSuite,
};

use crate::e2e::make_count_request;

async fn validate_search_lib_result(
    reader: Arc<ShardMirrorManager>,
    index_key: &IndexKey,
    shard_header: &ShardManifestHeader,
) -> Result<(), Box<dyn Error>> {
    // Good ranges: contains same table_id
    let ranges = vec![
        Range {
            start: shard_header.start_key.as_slice(),
            end: shard_header.end_key.as_slice(),
        },
        Range {
            start: shard_header.start_key.as_slice(),
            end: shard_header.end_key.as_slice(),
        },
    ];
    let res = make_count_request(reader.clone(), shard_header, index_key, ranges).await;
    assert!(res.is_ok());

    // Bad ranges: contains different table_id
    let base_start = table_start_key(index_key).as_slice().to_vec();
    let base_end = table_end_key(index_key).as_slice().to_vec();
    let other_index_key = IndexKey::new(
        index_key.keyspace_id,
        index_key.table_id + 100,
        index_key.index_id,
    );
    let other_start = table_start_key(&other_index_key).as_slice().to_vec();
    let other_end = table_end_key(&other_index_key).as_slice().to_vec();
    let ranges: Vec<Range<'_>> = vec![
        Range {
            start: base_start.as_slice(),
            end: base_end.as_slice(),
        },
        Range {
            start: other_start.as_slice(),
            end: other_end.as_slice(),
        },
    ];
    let res = make_count_request(reader.clone(), shard_header, index_key, ranges).await;
    assert!(
        res.is_err(),
        "count should fail with inconsistent table_id ranges"
    );

    // Bad range: empty range
    let ranges: Vec<Range<'_>> = vec![];
    let res = make_count_request(reader.clone(), shard_header, index_key, ranges).await;
    assert!(res.is_err(), "count should fail with empty ranges");

    Ok(())
}

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

    const READER_ADDR: &str = "127.0.0.1:19410";
    const READER_FLASH_ADDR: &str = "127.0.0.1:19411";

    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;

        // Add a reader node for testing
        let searchlib_config = make_default_reader_config(
            &test_suite.get_meta_service().config.load_full(),
            READER_ADDR.to_string(),
            READER_FLASH_ADDR.to_string(),
        );
        test_suite
            .add_real_reader_node_with_cfg(searchlib_config)
            .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 e2e 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 get_shard_local_cache_request works correctly
        let key_ranges = physical_table_ids
            .iter()
            .copied()
            .map(|physical_table_id| {
                let index_key = IndexKey::new(keyspace_id, physical_table_id, index_id);
                KeyRange {
                    start_key: table_start_key(&index_key).to_vec(),
                    end_key: table_end_key(&index_key).to_vec(),
                }
            })
            .collect::<Vec<_>>();
        let response = test_suite
            .send_get_shard_local_cache_request(&logical_index_key, key_ranges)
            .await
            .unwrap();
        test_suite
            .check_valid_shard_allocation(&response, 0)
            .unwrap();
        assert_eq!(
            response.shard_local_cache_infos.len(),
            physical_table_ids.len()
        );

        // Check if searchlib works correctly with partition table
        // Cause tidb planner transfer logical table_id tiflash then search_lib,
        // so we do the same.
        let shard_mirror_manager = test_suite
            .get_reader(READER_ADDR)
            .unwrap_right()
            .read()
            .await
            .get_service()
            .shard_mirror_manager
            .clone();
        for shard_info in response.shard_local_cache_infos.iter() {
            let shard_header = shard_info.shard.as_ref().expect("missing shard header");
            validate_search_lib_result(
                shard_mirror_manager.clone(),
                &logical_index_key,
                shard_header,
            )
            .await
            .unwrap();
        }

        // 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()
            );
        }

        test_suite.stop_readers().await.unwrap();
        test_suite.stop().await.unwrap();
    });

    Ok(())
}