// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.

use std::u64;

use futures::executor::block_on;
use kvproto::{
    coprocessor::{KeyRange, Request},
    kvrpcpb::{Context, IsolationLevel},
};
use protobuf::Message;
use test_coprocessor::*;
use tidb_query_common::storage::{
    scanner::{RangesScanner, RangesScannerOptions},
    Range,
};
use tikv::{
    coprocessor::{dag::TikvStorage, *},
    storage::{Engine, SnapshotStore},
};
use tipb::{ChecksumAlgorithm, ChecksumRequest, ChecksumResponse, ChecksumScanOn};
use txn_types::TimeStamp;

fn new_checksum_request(range: KeyRange, scan_on: ChecksumScanOn) -> Request {
    let mut ctx = Context::default();
    ctx.set_isolation_level(IsolationLevel::Si);

    let mut checksum = ChecksumRequest::default();
    checksum.set_scan_on(scan_on);
    checksum.set_algorithm(ChecksumAlgorithm::Crc64Xor);

    let mut req = Request::default();
    req.set_start_ts(u64::MAX);
    req.set_context(ctx);
    req.set_tp(REQ_TYPE_CHECKSUM);
    req.set_data(checksum.write_to_bytes().unwrap());
    req.mut_ranges().push(range);
    req
}

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

    let product = ProductTable::new();
    let (store, endpoint, _) = init_data_with_commit(&product, &data, true);

    for column in &[&product["id"], &product["name"], &product["count"]] {
        assert!(column.index >= 0);
        let (range, scan_on) = if column.index == 0 {
            let range = product.get_record_range_all();
            (range, ChecksumScanOn::Table)
        } else {
            let range = product.get_index_range_all(column.index);
            (range, ChecksumScanOn::Index)
        };
        let request = new_checksum_request(range.clone(), scan_on);
        let expected = reversed_checksum_crc64_xor(&store, range);

        let response = handle_request(&endpoint, request);
        let mut resp = ChecksumResponse::default();
        resp.merge_from_bytes(response.get_data()).unwrap();
        assert_eq!(resp.get_checksum(), expected);
        assert_eq!(resp.get_total_kvs(), data.len() as u64);
    }
}

fn reversed_checksum_crc64_xor<E: Engine>(store: &Store<E>, range: KeyRange) -> u64 {
    let store = SnapshotStore::new(
        store.get_engine().snapshot(Default::default()).unwrap(),
        TimeStamp::max(),
        IsolationLevel::Si,
        true,
        Default::default(),
        Default::default(),
        false,
    );
    let mut scanner = RangesScanner::new(RangesScannerOptions {
        storage: TikvStorage::new(store, false),
        ranges: vec![Range::from_pb_range(range, false)],
        scan_backward_in_range: true,
        is_key_only: false,
        is_scanned_range_aware: false,
    });

    let mut checksum = 0;
    let digest = crc64fast::Digest::new();
    while let Some((k, v)) = block_on(scanner.next()).unwrap() {
        let mut digest = digest.clone();
        digest.write(&k);
        digest.write(&v);
        checksum ^= digest.sum64();
    }
    checksum
}