use anyhow::Result;
use scopeguard::defer;
use tici::{
    config::{make_default_test_meta_config, make_default_test_worker_config},
    worker::mocks::make_frag_meta_extended,
};
use tici_common::{codec::table_raw_key, index::IndexKey};
use tici_grpc::{
    AppendBaseFragMetaRequest, ErrorCode, PreSplitImportShardMeta, PreSplitImportShardsRequest,
};
use tici_import::ImportJobMeta;
use tici_system::readable::ReadableSize;
use tici_test_utils::test_utils::{
    AppendExtraParam, mock_frag_path,
    test_suite::{E2ETestSuite, *},
};
use tidb_query_datatype::{
    codec::{Datum, datum, table},
    expr::EvalContext,
};

fn encode_import_fulltext_key(index_key: &IndexKey, values: &[Datum]) -> Vec<u8> {
    let encoded = datum::encode_key(&mut EvalContext::default(), values).unwrap();
    table::encode_index_seek_key(index_key.table_id, index_key.index_id, &encoded)
}

fn build_fulltext_split_key(index_key: &IndexKey, handle_values: &[Datum]) -> Vec<u8> {
    let mut suffix = b"_r".to_vec();
    suffix.extend(datum::encode_key(&mut EvalContext::default(), handle_values).unwrap());
    table_raw_key(index_key, &suffix)
}

#[test]
fn test_split_strategy_round_robin() -> Result<()> {
    run_split_strategy_test("round-robin", SplitStrategyExpectation::DifferentWriters)
}

#[test]
fn test_split_strategy_balanced_count() -> Result<()> {
    run_split_strategy_test(
        "balanced-count",
        SplitStrategyExpectation::SameWriterOnLessLoadedNode,
    )
}

enum SplitStrategyExpectation {
    DifferentWriters,
    SameWriterOnLessLoadedNode,
}

fn run_split_strategy_test(
    writer_scheduling_strategy: &str,
    expectation: SplitStrategyExpectation,
) -> Result<()> {
    let rt: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();

    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });
    let mut cfg = make_default_test_meta_config(false);
    cfg.shard.max_size = ReadableSize::mb(16);
    cfg.shard.writer_scheduling_strategy = writer_scheduling_strategy.to_string();
    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();

        test_suite
            .add_mocked_worker_node(Some("127.0.0.1:8556".to_string()))
            .await;
        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.check_meta_subscriber().await;

        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        test_suite.check_meta_subscriber().await;
        test_suite.validate_add_index_result(&index_key).await;

        let shard = test_suite.get_first_shard().await;
        let original_writer = shard.writer.clone();
        test_suite
            .must_append_delta_by_shard(
                &shard,
                "CCC",
                15,
                AppendExtraParam {
                    do_manifest_check: true,
                    upload_to_s3: false,
                    size: Some(cfg.shard.max_size.0),
                    count: Some(10000),
                    segs_count: Some(4),
                    ..Default::default()
                },
            )
            .await;

        must_finally(async || {
            // Because split can only triggered by append event, so the split can only take
            // place once.
            let shards = test_suite.collect_shards_from_restore().await;
            if shards.len() != 2 {
                return false;
            }
            match expectation {
                SplitStrategyExpectation::DifferentWriters => {
                    shards.first().unwrap().writer != shards.get(1).unwrap().writer
                }
                SplitStrategyExpectation::SameWriterOnLessLoadedNode => {
                    let split_writer = &shards.first().unwrap().writer;
                    split_writer == &shards.get(1).unwrap().writer
                        && split_writer != &original_writer
                }
            }
        })
        .await;

        test_suite.stop_worker().await.unwrap();

        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });
    Ok(())
}

// TODO(P1) Fix this test and its assertion.
// #[test]
fn test_append_base_frag_does_not_auto_split_during_active_import_job() -> Result<()> {
    let rt: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();

    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(false);
    cfg.shard.auto_compaction = false;
    cfg.shard.enable_split = true;
    cfg.shard.enable_merge = false;
    cfg.shard.max_size = ReadableSize::mb(16);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();
        test_suite
            .add_mocked_worker_node(Some("127.0.0.1:8556".to_string()))
            .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.check_meta_subscriber().await;
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        test_suite.validate_add_index_result(&index_key).await;

        let shard = test_suite.get_first_shard().await;
        let shard_id = shard.shard_id;
        let worker_addr = test_suite.get_one_worker_node_addr();

        let job_id = test_suite
            .get_meta_service()
            .meta_store
            .create_or_get_import_job(&ImportJobMeta {
                tidb_task_id: "append-base-split-blocked-by-import".to_string(),
                keyspace_id: index_key.keyspace_id,
                table_id: index_key.table_id,
                index_ids: vec![index_key.index_id],
                job_id: 0,
            })
            .await
            .unwrap();

        test_suite
            .must_append_base_frag(
                &worker_addr,
                shard_id,
                &index_key,
                vec![make_frag_meta_extended(
                    &mock_frag_path("import-base-over-limit"),
                    &shard.range,
                    cfg.shard.max_size.0 * 2,
                    1,
                    1,
                )],
            )
            .await
            .unwrap();

        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        assert_eq!(test_suite.collect_shards_from_restore().await.len(), 1);

        test_suite
            .get_meta_service()
            .meta_store
            .finish_import_job(&[job_id])
            .await
            .unwrap();

        test_suite.stop_worker().await.unwrap();
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_presplit_import_shards_via_grpc() -> Result<()> {
    let rt: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();

    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(false);
    cfg.shard.enable_merge = false;
    cfg.shard.max_size = ReadableSize::mb(16);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();
        test_suite
            .add_mocked_worker_node(Some("127.0.0.1:8556".to_string()))
            .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.check_meta_subscriber().await;
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        test_suite.validate_add_index_result(&index_key).await;

        let import_key_0 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"debug".to_vec()),
                Datum::Bytes(b"doc-0".to_vec()),
                Datum::I64(1),
                Datum::I64(10),
            ],
        );
        let import_key_1 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"error".to_vec()),
                Datum::Bytes(b"doc-1".to_vec()),
                Datum::I64(2),
                Datum::I64(20),
            ],
        );
        let import_key_2 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"fatal".to_vec()),
                Datum::Bytes(b"doc-2".to_vec()),
                Datum::I64(3),
                Datum::I64(30),
            ],
        );
        let import_key_3 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"panic".to_vec()),
                Datum::Bytes(b"doc-3".to_vec()),
                Datum::I64(4),
                Datum::I64(40),
            ],
        );
        let mut client = test_suite.get_meta_client();
        let response = client
            .pre_split_import_shards(PreSplitImportShardsRequest {
                tidb_task_id: "presplit-import-task".to_string(),
                table_id: index_key.table_id,
                index_ids: vec![index_key.index_id],
                start_key: import_key_0.clone(),
                end_key: import_key_3.clone(),
                total_kv_size: cfg.shard.max_size.0 * 3,
                total_kv_cnt: 300,
                meta_groups: vec![
                    PreSplitImportShardMeta {
                        ele_id: 1,
                        start_key: import_key_0.clone(),
                        end_key: import_key_1.clone(),
                        total_kv_size: cfg.shard.max_size.0,
                        total_kv_cnt: 100,
                        data_file_count: 1,
                        stat_file_count: 1,
                    },
                    PreSplitImportShardMeta {
                        ele_id: 2,
                        start_key: import_key_1.clone(),
                        end_key: import_key_2.clone(),
                        total_kv_size: cfg.shard.max_size.0,
                        total_kv_cnt: 100,
                        data_file_count: 1,
                        stat_file_count: 1,
                    },
                    PreSplitImportShardMeta {
                        ele_id: 3,
                        start_key: import_key_2.clone(),
                        end_key: import_key_3.clone(),
                        total_kv_size: cfg.shard.max_size.0,
                        total_kv_cnt: 100,
                        data_file_count: 1,
                        stat_file_count: 1,
                    },
                ],
                keyspace_id: index_key.keyspace_id,
                ..Default::default()
            })
            .await
            .unwrap()
            .into_inner();

        assert_eq!(response.status, ErrorCode::Success as i32);
        assert_eq!(response.split_keys, vec![import_key_1, import_key_2]);
        assert_eq!(response.index_results.len(), 1);
        assert_eq!(response.index_results[0].split_count, 2);

        must_finally(async || test_suite.collect_shards_from_restore().await.len() == 3).await;

        test_suite.stop_worker().await.unwrap();
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_presplit_import_shards_waits_for_shard_task_serialization() -> Result<()> {
    let rt: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();

    defer!({
        fail::remove("worker_service_mock_meta_service_url");
        fail::remove("append_base_frag_meta_before_get_shard");
    });

    let mut cfg = make_default_test_meta_config(false);
    cfg.shard.enable_merge = false;
    cfg.shard.max_size = ReadableSize::mb(16);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();
        test_suite
            .add_mocked_worker_node(Some("127.0.0.1:8556".to_string()))
            .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.check_meta_subscriber().await;
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        test_suite.validate_add_index_result(&index_key).await;

        let shard = test_suite.get_first_shard().await;
        let import_key_0 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"debug".to_vec()),
                Datum::Bytes(b"doc-0".to_vec()),
                Datum::I64(1),
                Datum::I64(10),
            ],
        );
        let import_key_1 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"error".to_vec()),
                Datum::Bytes(b"doc-1".to_vec()),
                Datum::I64(2),
                Datum::I64(20),
            ],
        );
        let import_key_2 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"fatal".to_vec()),
                Datum::Bytes(b"doc-2".to_vec()),
                Datum::I64(3),
                Datum::I64(30),
            ],
        );

        fail::cfg("append_base_frag_meta_before_get_shard", "pause").unwrap();

        let mut append_client = test_suite.get_meta_internal_client();
        let append_req = AppendBaseFragMetaRequest {
            keyspace_id: index_key.keyspace_id,
            table_id: index_key.table_id,
            index_id: index_key.index_id,
            frag_metas: vec![make_frag_meta_extended(
                &mock_frag_path("presplit-serialize-base"),
                &shard.range,
                cfg.shard.max_size.0 / 4,
                1,
                1,
            )],
        };
        let append_handle =
            tokio::spawn(async move { append_client.append_base_frag_meta(append_req).await });

        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        assert!(
            !append_handle.is_finished(),
            "append request finished before presplit was issued"
        );

        let mut presplit_client = test_suite.get_meta_client();
        let presplit_handle = tokio::spawn(async move {
            presplit_client
                .pre_split_import_shards(PreSplitImportShardsRequest {
                    tidb_task_id: "presplit-serialization".to_string(),
                    table_id: index_key.table_id,
                    index_ids: vec![index_key.index_id],
                    start_key: import_key_0.clone(),
                    end_key: import_key_2.clone(),
                    total_kv_size: cfg.shard.max_size.0 * 2,
                    total_kv_cnt: 200,
                    meta_groups: vec![
                        PreSplitImportShardMeta {
                            ele_id: 1,
                            start_key: import_key_0,
                            end_key: import_key_1.clone(),
                            total_kv_size: cfg.shard.max_size.0,
                            total_kv_cnt: 100,
                            data_file_count: 1,
                            stat_file_count: 1,
                        },
                        PreSplitImportShardMeta {
                            ele_id: 2,
                            start_key: import_key_1,
                            end_key: import_key_2,
                            total_kv_size: cfg.shard.max_size.0,
                            total_kv_cnt: 100,
                            data_file_count: 1,
                            stat_file_count: 1,
                        },
                    ],
                    keyspace_id: index_key.keyspace_id,
                    ..Default::default()
                })
                .await
        });

        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
        assert!(
            !presplit_handle.is_finished(),
            "presplit request bypassed shard task serialization"
        );

        fail::remove("append_base_frag_meta_before_get_shard");

        let append_resp = append_handle.await.unwrap().unwrap().into_inner();
        assert_eq!(append_resp.status, ErrorCode::Success as i32);

        let presplit_resp = presplit_handle.await.unwrap().unwrap().into_inner();
        assert_eq!(presplit_resp.status, ErrorCode::Success as i32);
        assert_eq!(presplit_resp.index_results.len(), 1);
        assert_eq!(presplit_resp.index_results[0].split_count, 1);

        must_finally(async || test_suite.collect_shards_from_restore().await.len() == 2).await;

        test_suite.stop_worker().await.unwrap();
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}

#[test]
fn test_presplit_import_shards_accepts_parser_shaped_fulltext_keys() -> Result<()> {
    let rt: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()?;
    fail::cfg("worker_service_mock_meta_service_url", "return").unwrap();

    defer!({
        fail::remove("worker_service_mock_meta_service_url");
    });

    let mut cfg = make_default_test_meta_config(false);
    cfg.shard.enable_merge = false;
    cfg.shard.max_size = ReadableSize::mb(16);

    rt.block_on(async {
        let mut test_suite =
            E2ETestSuite::new_with_mock_cfg(cfg.clone(), make_default_test_worker_config(true))
                .await
                .unwrap();
        test_suite
            .add_mocked_worker_node(Some("127.0.0.1:8556".to_string()))
            .await;

        let index_key = IndexKey::new_nullspace(1, 1001);
        test_suite.must_create_index(&index_key).await.unwrap();
        test_suite.check_meta_subscriber().await;
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
        test_suite.validate_add_index_result(&index_key).await;

        // Import fulltext keys can contain only parser output plus handle columns,
        // even when the index metadata declares multiple indexed columns.
        let import_key_0 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"debug".to_vec()),
                Datum::I64(1),
                Datum::I64(10),
            ],
        );
        let import_key_1 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"error".to_vec()),
                Datum::I64(2),
                Datum::I64(20),
            ],
        );
        let import_key_2 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"fatal".to_vec()),
                Datum::I64(3),
                Datum::I64(30),
            ],
        );
        let import_key_3 = encode_import_fulltext_key(
            &index_key,
            &[
                Datum::Bytes(b"panic".to_vec()),
                Datum::I64(4),
                Datum::I64(40),
            ],
        );
        let mut client = test_suite.get_meta_client();
        let response = client
            .pre_split_import_shards(PreSplitImportShardsRequest {
                tidb_task_id: "presplit-import-parser-keys".to_string(),
                table_id: index_key.table_id,
                index_ids: vec![index_key.index_id],
                start_key: import_key_0.clone(),
                end_key: import_key_3.clone(),
                total_kv_size: cfg.shard.max_size.0 * 3,
                total_kv_cnt: 300,
                meta_groups: vec![
                    PreSplitImportShardMeta {
                        ele_id: 1,
                        start_key: import_key_0.clone(),
                        end_key: import_key_1.clone(),
                        total_kv_size: cfg.shard.max_size.0,
                        total_kv_cnt: 100,
                        data_file_count: 1,
                        stat_file_count: 1,
                    },
                    PreSplitImportShardMeta {
                        ele_id: 2,
                        start_key: import_key_1.clone(),
                        end_key: import_key_2.clone(),
                        total_kv_size: cfg.shard.max_size.0,
                        total_kv_cnt: 100,
                        data_file_count: 1,
                        stat_file_count: 1,
                    },
                    PreSplitImportShardMeta {
                        ele_id: 3,
                        start_key: import_key_2.clone(),
                        end_key: import_key_3.clone(),
                        total_kv_size: cfg.shard.max_size.0,
                        total_kv_cnt: 100,
                        data_file_count: 1,
                        stat_file_count: 1,
                    },
                ],
                keyspace_id: index_key.keyspace_id,
                ..Default::default()
            })
            .await
            .unwrap()
            .into_inner();

        assert_eq!(response.status, ErrorCode::Success as i32);
        assert_eq!(response.split_keys, vec![import_key_1, import_key_2]);
        assert_eq!(response.index_results.len(), 1);
        assert_eq!(response.index_results[0].split_count, 2);

        must_finally(async || test_suite.collect_shards_from_restore().await.len() == 3).await;

        test_suite.stop_worker().await.unwrap();
        test_suite.check_meta_subscriber().await;
        test_suite.stop_meta().await.unwrap();
    });

    Ok(())
}