// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use kvproto::kvrpcpb::{Context, IsolationLevel};
use more_asserts::{assert_ge, assert_le};
use protobuf::Message;
use test_coprocessor::*;
use test_storage::*;
use tidb_query_datatype::{
    FieldTypeTp,
    codec::{Datum, datum},
    expr::EvalContext,
};
use tipb::{ScalarFuncSig, SelectResponse};
use tipb_helper::ExprDefBuilder;

#[test]
fn test_deadline() {
    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &[]);
    let req = DagSelect::from(&product).build();

    fail::cfg("deadline_check_fail", "return()").unwrap();
    let resp = handle_request(&endpoint, req);
    let region_err = resp.get_region_error();
    assert_eq!(
        region_err.get_server_is_busy().reason,
        "deadline is exceeded".to_string()
    );
    assert_eq!(
        region_err.get_message(),
        "Coprocessor task terminated due to exceeding the deadline"
    );
}

#[test]
fn test_deadline_2() {
    // It should not even take any snapshots when request is outdated from the
    // beginning.
    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &[]);
    let req = DagSelect::from(&product).build();

    fail::cfg("rockskv_async_snapshot", "panic").unwrap();
    fail::cfg("deadline_check_fail", "return()").unwrap();
    let resp = handle_request(&endpoint, req);
    let region_err = resp.get_region_error();
    assert_eq!(
        region_err.get_server_is_busy().reason,
        "deadline is exceeded".to_string()
    );
    assert_eq!(
        region_err.get_message(),
        "Coprocessor task terminated due to exceeding the deadline"
    );
}

/// Test deadline exceeded when request is handling
/// Note: only
#[test]
fn test_deadline_3() {
    let data = vec![
        (1, Some("name:0"), 2),
        (2, Some("name:4"), 3),
        (4, Some("name:3"), 1),
        (5, Some("name:1"), 4),
    ];

    let product = ProductTable::new();
    let (_, endpoint) = {
        let engine = tikv::storage::TestEngineBuilder::new().build().unwrap();
        let cfg = tikv::server::Config {
            end_point_request_max_handle_duration: tikv_util::config::ReadableDuration::secs(1),
            ..Default::default()
        };
        init_data_with_details(Context::default(), engine, &product, &data, true, &cfg)
    };
    let req = DagSelect::from(&product).build();

    fail::cfg("kv_cursor_seek", "sleep(2000)").unwrap();
    fail::cfg("copr_batch_initial_size", "return(1)").unwrap();
    let cop_resp = handle_request(&endpoint, req);
    let mut resp = SelectResponse::default();
    resp.merge_from_bytes(cop_resp.get_data()).unwrap();

    let region_err = cop_resp.get_region_error();
    assert_eq!(
        region_err.get_server_is_busy().reason,
        "deadline is exceeded".to_string()
    );
    assert_eq!(
        region_err.get_message(),
        "Coprocessor task terminated due to exceeding the deadline"
    );
}

#[test]
fn test_parse_request_failed() {
    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &[]);
    let req = DagSelect::from(&product).build();

    fail::cfg("coprocessor_parse_request", "return()").unwrap();
    let resp = handle_request(&endpoint, req);

    assert!(resp.get_other_error().contains("unsupported tp"));
}

#[test]
fn test_parse_request_failed_2() {
    // It should not even take any snapshots when parse failed.
    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &[]);
    let req = DagSelect::from(&product).build();

    fail::cfg("rockskv_async_snapshot", "panic").unwrap();
    fail::cfg("coprocessor_parse_request", "return()").unwrap();
    let resp = handle_request(&endpoint, req);

    assert!(resp.get_other_error().contains("unsupported tp"));
}

#[test]
fn test_readpool_full() {
    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &[]);
    let req = DagSelect::from(&product).build();

    fail::cfg("future_pool_spawn_full", "return()").unwrap();
    let resp = handle_request(&endpoint, req);

    assert!(resp.get_region_error().has_server_is_busy());
}

#[test]
fn test_snapshot_failed() {
    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &[]);
    let req = DagSelect::from(&product).build();

    fail::cfg("rockskv_async_snapshot", "return()").unwrap();
    let resp = handle_request(&endpoint, req);

    assert!(resp.get_other_error().contains("snapshot failed"));
}

#[test]
fn test_snapshot_failed_2() {
    let product = ProductTable::new();
    let (store, endpoint) = init_with_data(&product, &[]);
    let req = DagSelect::from(&product).build();

    store.get_engine().trigger_not_leader();
    let resp = handle_request(&endpoint, req);

    assert!(resp.get_region_error().has_not_leader());
}

#[test]
fn test_storage_error() {
    let data = vec![(1, Some("name:0"), 2), (2, Some("name:4"), 3)];

    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &data);
    let req = DagSelect::from(&product).build();

    fail::cfg("kv_cursor_seek", "return()").unwrap();
    let resp = handle_request(&endpoint, req);

    assert!(resp.get_other_error().contains("kv cursor seek error"));
}

#[test]
fn test_region_error_in_scan() {
    let data = vec![
        (1, Some("name:0"), 2),
        (2, Some("name:4"), 3),
        (4, Some("name:3"), 1),
        (5, Some("name:1"), 4),
    ];

    let product = ProductTable::new();
    let (_cluster, raft_engine, mut ctx) = new_raft_engine(1, "");
    ctx.set_isolation_level(IsolationLevel::Si);

    let (_, endpoint) =
        init_data_with_engine_and_commit(ctx.clone(), raft_engine, &product, &data, true);

    fail::cfg("region_snapshot_seek", "return()").unwrap();
    let req = DagSelect::from(&product).build_with(ctx, &[0]);
    let resp = handle_request(&endpoint, req);

    assert!(
        resp.get_region_error()
            .get_message()
            .contains("region seek error")
    );
}

#[test]
fn test_paging_scan() {
    let data = vec![
        (1, Some("name:0"), 2),
        (2, Some("name:4"), 3),
        (4, Some("name:3"), 1),
        (5, Some("name:1"), 4),
    ];

    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &data);
    // set batch size and grow size to 1, so that only 1 row will be scanned in each
    // batch.
    fail::cfg("copr_batch_initial_size", "return(1)").unwrap();
    fail::cfg("copr_batch_grow_size", "return(1)").unwrap();
    for desc in [false, true] {
        for paging_size in 1..=4 {
            let mut exp = data.clone();
            if desc {
                exp.reverse();
            }

            let req = DagSelect::from(&product)
                .paging_size(paging_size as u64)
                .desc(desc)
                .build();
            let resp = handle_request(&endpoint, req);
            let mut select_resp = SelectResponse::default();
            select_resp.merge_from_bytes(resp.get_data()).unwrap();

            let mut row_count = 0;
            let spliter = DagChunkSpliter::new(select_resp.take_chunks().into(), 3);
            for (row, (id, name, cnt)) in spliter.zip(exp) {
                let name_datum = name.unwrap().as_bytes().into();
                let expected_encoded = datum::encode_value(
                    &mut EvalContext::default(),
                    &[Datum::I64(id), name_datum, Datum::I64(cnt)],
                )
                .unwrap();
                let result_encoded =
                    datum::encode_value(&mut EvalContext::default(), &row).unwrap();
                assert_eq!(result_encoded, &*expected_encoded);
                row_count += 1;
            }
            assert_eq!(row_count, paging_size);

            let res_range = resp.get_range();
            let (res_start_key, res_end_key) = match desc {
                true => (res_range.get_end(), res_range.get_start()),
                false => (res_range.get_start(), res_range.get_end()),
            };
            let start_key = match desc {
                true => product.get_record_range_one(i64::MAX),
                false => product.get_record_range_one(i64::MIN),
            };
            let end_key = match desc {
                true => product.get_record_range_one(data[data.len() - paging_size].0),
                false => product.get_record_range_one(data[paging_size - 1].0),
            };
            assert_eq!(res_start_key, start_key.get_start());
            assert_ge!(res_end_key, end_key.get_start());
            assert_le!(res_end_key, end_key.get_end());
        }

        // test limit with early return
        let req = DagSelect::from(&product)
            .paging_size(2)
            .limit(1)
            .desc(desc)
            .build();
        let resp = handle_request(&endpoint, req);
        assert!(resp.range.is_none());
        assert!(resp.range.is_none());

        let agg_req = DagSelect::from(&product)
            .count(&product["count"])
            .group_by(&[&product["name"]])
            .output_offsets(Some(vec![0, 1]))
            .desc(desc)
            .paging_size(2)
            .build();
        let resp = handle_request(&endpoint, agg_req);
        assert!(resp.range.is_some());
    }
}

#[test]
fn test_paging_scan_bytes() {
    let data = vec![
        (1, Some("name:0"), 2),
        (2, Some("name:4"), 3),
        (4, Some("name:3"), 1),
        (5, Some("name:1"), 4),
    ];

    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &data);
    // Force one row per batch so the byte budget is evaluated after every row,
    // making the truncation point deterministic without depending on the exact
    // per-row encoded size.
    fail::cfg("copr_batch_initial_size", "return(1)").unwrap();
    fail::cfg("copr_batch_grow_size", "return(1)").unwrap();

    // Plain table scan, both directions, with the precise resume-range check. A
    // 1-byte budget is exceeded after the first row's MVCC processed_size
    // (encoded key + value, always > 1 byte), so the scan stops after exactly
    // one row and hands back a resume range that starts at the table boundary and
    // stops within that single row's key range. An unbounded budget drains
    // naturally and returns no resume range.
    for desc in [false, true] {
        let req = DagSelect::from(&product)
            .paging_size_bytes(1)
            .desc(desc)
            .build();
        let resp = handle_request(&endpoint, req);
        let mut select_resp = SelectResponse::default();
        select_resp.merge_from_bytes(resp.get_data()).unwrap();
        let row_count = DagChunkSpliter::new(select_resp.take_chunks().into(), 3).count();
        assert_eq!(row_count, 1, "byte budget must truncate to a single row");

        let res_range = resp.get_range();
        let (res_start_key, res_end_key) = match desc {
            true => (res_range.get_end(), res_range.get_start()),
            false => (res_range.get_start(), res_range.get_end()),
        };
        let start_key = match desc {
            true => product.get_record_range_one(i64::MAX),
            false => product.get_record_range_one(i64::MIN),
        };
        let stop_id = match desc {
            true => data[data.len() - 1].0,
            false => data[0].0,
        };
        let end_key = product.get_record_range_one(stop_id);
        assert_eq!(res_start_key, start_key.get_start());
        assert_ge!(res_end_key, end_key.get_start());
        assert_le!(res_end_key, end_key.get_end());

        let req = DagSelect::from(&product)
            .paging_size_bytes(u64::MAX)
            .desc(desc)
            .build();
        let resp = handle_request(&endpoint, req);
        let mut select_resp = SelectResponse::default();
        select_resp.merge_from_bytes(resp.get_data()).unwrap();
        let row_count = DagChunkSpliter::new(select_resp.take_chunks().into(), 3).count();
        assert_eq!(
            row_count,
            data.len(),
            "ample byte budget must return all rows"
        );
        assert!(
            resp.range.is_none(),
            "a draining scan must not set a resume range"
        );
    }

    // Selection above the scan, the subtle case: the byte budget is charged on
    // *storage-scanned* bytes, below filtering, so a page that emits zero rows
    // still stops -- and its resume range must advance past the scanned row, or
    // the caller re-scans it forever. `id > 1` drops the first scanned row (asc
    // handle order starts at id = 1), so a 1-byte budget stops after scanning it
    // and resumes with zero rows; the resume range must still start at the table
    // boundary and stop within that first row's key range. An unbounded budget
    // keeps the three matching rows and drains with no range.
    let selection = |bytes: u64| {
        DagSelect::from(&product)
            .where_expr(
                ExprDefBuilder::scalar_func(ScalarFuncSig::GtInt, FieldTypeTp::LongLong)
                    .push_child(ExprDefBuilder::column_ref(0, FieldTypeTp::LongLong))
                    .push_child(ExprDefBuilder::constant_int(1))
                    .build(),
            )
            .paging_size_bytes(bytes)
            .build()
    };
    let resp = handle_request(&endpoint, selection(1));
    let mut select_resp = SelectResponse::default();
    select_resp.merge_from_bytes(resp.get_data()).unwrap();
    assert_eq!(
        DagChunkSpliter::new(select_resp.take_chunks().into(), 3).count(),
        0,
        "a fully-filtered page must emit no rows"
    );
    let res_range = resp.get_range();
    let first_row = product.get_record_range_one(data[0].0);
    assert_eq!(
        res_range.get_start(),
        product.get_record_range_one(i64::MIN).get_start()
    );
    assert_ge!(res_range.get_end(), first_row.get_start());
    assert_le!(res_range.get_end(), first_row.get_end());

    let resp = handle_request(&endpoint, selection(u64::MAX));
    let mut select_resp = SelectResponse::default();
    select_resp.merge_from_bytes(resp.get_data()).unwrap();
    assert_eq!(
        DagChunkSpliter::new(select_resp.take_chunks().into(), 3).count(),
        3,
        "an ample budget must return every matching row"
    );
    assert!(resp.range.is_none(), "a drained scan must not resume");

    // The remaining gate-allowed executors: an IndexScan leaf and a Projection
    // pass-through. They must forward `peek_scanned_bytes_sum` so a 1-byte budget
    // truncates to the first scanned row with a resume range, and drain with none
    // under an unbounded budget. The resume-range *boundary* is already pinned by
    // the table-scan and selection cases above, which share this RangesScanner.
    let check = |make: &dyn Fn() -> DagSelect, col_cnt: usize, label: &str| {
        let resp = handle_request(&endpoint, make().paging_size_bytes(1).build());
        let mut select_resp = SelectResponse::default();
        select_resp.merge_from_bytes(resp.get_data()).unwrap();
        let row_count = DagChunkSpliter::new(select_resp.take_chunks().into(), col_cnt).count();
        assert_eq!(
            row_count, 1,
            "{}: 1-byte budget must truncate to one row",
            label
        );
        assert!(
            resp.range.is_some(),
            "{}: a truncated scan must resume",
            label
        );

        let resp = handle_request(&endpoint, make().paging_size_bytes(u64::MAX).build());
        let mut select_resp = SelectResponse::default();
        select_resp.merge_from_bytes(resp.get_data()).unwrap();
        let row_count = DagChunkSpliter::new(select_resp.take_chunks().into(), col_cnt).count();
        assert_eq!(
            row_count,
            data.len(),
            "{}: ample budget returns all rows",
            label
        );
        assert!(
            resp.range.is_none(),
            "{}: a drained scan must not resume",
            label
        );
    };
    check(
        &|| DagSelect::from_index(&product, &product["name"]),
        3,
        "index scan",
    );
    check(
        &|| {
            DagSelect::from(&product)
                .projection(vec![
                    ExprDefBuilder::column_ref(0, FieldTypeTp::LongLong).build(),
                ])
                .output_offsets(Some(vec![0]))
        },
        1,
        "projection",
    );
}

#[test]
fn test_paging_scan_multi_ranges() {
    let data = vec![
        (1, Some("name:0"), 2),
        (2, Some("name:4"), 3),
        (3, Some("name:5"), 5),
        (4, Some("name:3"), 1),
        (5, Some("name:1"), 4),
    ];
    let product = ProductTable::new();
    let (_, endpoint) = init_with_data(&product, &data);
    // set batch size and grow size to 1, so that only 1 row will be scanned in each
    // batch.
    fail::cfg("copr_batch_initial_size", "return(1)").unwrap();
    fail::cfg("copr_batch_grow_size", "return(1)").unwrap();

    // test multi ranges with gap
    for desc in [true, false] {
        for paging_size in [3, 5] {
            let mut exp = [data[0], data[1], data[3], data[4]];
            if desc {
                exp.reverse();
            }

            let builder = DagSelect::from(&product)
                .paging_size(paging_size)
                .desc(desc);
            let mut range1 = builder.key_ranges[0].clone();
            range1.set_end(product.get_record_range_one(data[1].0).get_end().into());
            let mut range2 = builder.key_ranges[0].clone();
            range2.set_start(product.get_record_range_one(data[3].0).get_start().into());
            let key_ranges = vec![range1.clone(), range2.clone()];

            let req = builder.key_ranges(key_ranges).build();
            let resp = handle_request(&endpoint, req);
            let mut select_resp = SelectResponse::default();
            select_resp.merge_from_bytes(resp.get_data()).unwrap();

            let mut row_count = 0;
            let spliter = DagChunkSpliter::new(select_resp.take_chunks().into(), 3);
            for (row, (id, name, cnt)) in spliter.zip(exp) {
                let name_datum = name.unwrap().as_bytes().into();
                let expected_encoded = datum::encode_value(
                    &mut EvalContext::default(),
                    &[Datum::I64(id), name_datum, Datum::I64(cnt)],
                )
                .unwrap();
                let result_encoded =
                    datum::encode_value(&mut EvalContext::default(), &row).unwrap();
                assert_eq!(result_encoded, &*expected_encoded);
                row_count += 1;
            }
            let exp_len = if paging_size <= 4 {
                paging_size
            } else {
                exp.len() as u64
            };
            assert_eq!(row_count, exp_len);

            let res_range = resp.get_range();

            let (res_start_key, res_end_key) = match desc {
                true => (res_range.get_end(), res_range.get_start()),
                false => (res_range.get_start(), res_range.get_end()),
            };
            if paging_size != 5 {
                let start_key = match desc {
                    true => range2.get_end(),
                    false => range1.get_start(),
                };
                let end_id = match desc {
                    true => data[1].0,
                    false => data[3].0,
                };
                let end_key = product.get_record_range_one(end_id);
                assert_eq!(res_start_key, start_key);
                assert_ge!(res_end_key, end_key.get_start());
                assert_le!(res_end_key, end_key.get_end());
            } else {
                // drained.
                assert!(res_start_key.is_empty());
                assert!(res_end_key.is_empty());
            }
        }
    }
}