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

use std::{assert_matches::assert_matches, path::PathBuf, sync::Arc, time::Duration};

use bytes::{Buf, Bytes};
use kvengine::{
    dfs::{FileType, S3Fs},
    ia::{
        ia_file::{table_meta_file_local_path, IaFile},
        manager::IaManager,
        types::{FileSegmentData, FileSegmentIdent},
        util::{
            test_util::verify_local_segments, IaCapacity, IaManagerOptionsBuilder, LocalFileStore,
            LocalStore,
        },
    },
    table::{file::InMemFile, sstable, ChecksumType, InnerKey, Value, NO_COMPRESSION},
};
use proptest::prelude::*;
use rand::prelude::*;
use rstest::rstest;
use test_cloud_server::oss::prepare_dfs;
use test_util::init_log_for_test;
use tikv_util::{debug, info};

const BLOCK_SIZE: usize = 32;
const SEGMENT_SIZE: i64 = 64;
const FREQ_UPDATE_INTERVAL: Duration = Duration::from_secs(1);

prop_compose! {
    fn arb_range_args(min: u64, max: u64)
        (start in min..max)
        (
            start in Just(start),
            end in start+1..=max,
        )
        -> (u64, u64)
    {
        (start, end)
    }
}

#[rstest]
#[case(IaCapacity::MemoryAndDiskCap(300, PathBuf::from("ia"), 3000))]
#[case::memory(IaCapacity::MemoryCap(3000))]
#[case::big_cap(IaCapacity::MemoryAndDiskCap(1000, PathBuf::from("ia"), 10000))]
#[case::small_cap(IaCapacity::MemoryAndDiskCap(128, PathBuf::from("ia"), 1024))]
#[ignore]
fn test_read(#[case] mut ia_cap: IaCapacity) {
    init_log_for_test();

    let (temp_dir, mut oss, dfs_conf) = prepare_dfs("test");
    let temp_dir = temp_dir.path();

    let s3fs = S3Fs::new_from_config(dfs_conf);
    let _s3fs = s3fs.clone();

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()
        .unwrap();

    ia_cap.set_parent_dir(temp_dir.to_path_buf());
    let options = IaManagerOptionsBuilder::default()
        .capacity(ia_cap)
        .segment_size(SEGMENT_SIZE)
        .freq_update_interval(Duration::ZERO)
        .build()
        .unwrap();
    let (small_cap, main_cap) = (options.small_queue.cap, options.main_queue.cap);

    let rt = runtime.handle().clone();
    let (mgr, user_data, ia_file) = runtime.block_on(async move {
        let file_id = 42;
        let file_type = FileType::Sst;
        // user_data: About 6.8 KiB.
        let (file_data, user_data, table_meta_off) =
            make_sstable(file_id, BLOCK_SIZE, 200, 7, 5, thread_rng().gen_ratio(1, 2));
        info!("make sstable"; "file size" => file_data.len(), "user data size" => user_data.len());

        s3fs.put_object(
            s3fs.file_key(file_id, file_type),
            file_data,
            format!("{}.{}", file_id, file_type.suffix()),
        )
        .await
        .unwrap();

        let mgr = IaManager::new(options, s3fs.clone(), rt).await.unwrap();
        let table_meta_data =
            IaFile::prepare_table_meta(file_id, file_type, table_meta_off, temp_dir, &s3fs)
                .await
                .unwrap();
        let table_meta_file = InMemFile::new(file_id, table_meta_data);
        let ia_file =
            IaFile::open(file_id, file_type, Arc::new(table_meta_file), mgr.clone()).unwrap();

        let seg = ia_file.multi_read_async(0, user_data.len()).await.unwrap();
        assert_eq!(seg, user_data);

        let mut buf = vec![0; user_data.len()];
        ia_file.multi_read_at_async(&mut buf, 0).await.unwrap();
        assert_eq!(buf, user_data.chunk());

        (mgr, user_data, ia_file)
    });

    proptest!(|(
        (start_off, end_off) in arb_range_args(0, user_data.len() as u64)
    )| {
        debug!("test_read: start_off: {}, end_off: {}", start_off, end_off);
        let expected = user_data.slice(start_off as usize..end_off as usize);

        let seg = runtime.block_on(ia_file.multi_read_async(start_off, (end_off-start_off) as usize)).unwrap();
        prop_assert_eq!(&seg, &expected);

        let mut buf = vec![0; (end_off-start_off) as usize];
        runtime.block_on(ia_file.multi_read_at_async(&mut buf, start_off)).unwrap();
        prop_assert_eq!(buf, expected.chunk());
    });

    runtime.block_on(async {
        mgr.flush_tasks(Duration::from_secs(5)).await.unwrap();

        let segments = mgr.get_local_segments().await;
        verify_local_segments(&segments, small_cap, main_cap, Some(user_data.len() as u64));
    });

    info!("cache hit rate: {}", mgr.cache_hit_rate());
    oss.shutdown();
}

#[test]
fn test_init() {
    init_log_for_test();

    let (temp_dir, mut oss, dfs_conf) = prepare_dfs("test");
    let temp_dir = temp_dir.path();

    let s3fs = S3Fs::new_from_config(dfs_conf);
    let _s3fs = s3fs.clone();

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()
        .unwrap();
    let rt = runtime.handle().clone();
    runtime.block_on(async move {
        let local_path = temp_dir.join("ia");
        let ia_cap = IaCapacity::MemoryAndDiskCap(0, local_path.clone(), 1000);
        let options = IaManagerOptionsBuilder::default()
            .capacity(ia_cap)
            .segment_size(SEGMENT_SIZE)
            .freq_update_interval(FREQ_UPDATE_INTERVAL)
            .build()
            .unwrap();

        let file_type = FileType::Sst;

        {
            let mgr = IaManager::new(options.clone(), s3fs.clone(), rt.clone())
                .await
                .unwrap();

            let mut files = Vec::with_capacity(10);
            for i in 1..10 {
                let file_id = i as u64;
                let (file_data, _, table_meta_off) =
                    make_sstable(file_id, BLOCK_SIZE, 10, 7, 5, false);

                s3fs.put_object(
                    s3fs.file_key(file_id, file_type),
                    file_data,
                    format!("{}.{}", file_id, file_type.suffix()),
                )
                .await
                .unwrap();

                files.push((file_id, file_type, table_meta_off));
            }

            for (file_id, file_type, table_meta_off) in files {
                IaFile::prepare_table_meta(file_id, file_type, table_meta_off, &local_path, &s3fs)
                    .await
                    .unwrap();
            }

            let ia1 = IaFile::open_in_path(1, file_type, &local_path, mgr.clone()).unwrap();
            let _ = ia1.multi_read_async(5, 10).await.unwrap();
            let _ = ia1.multi_read_async(5, 10).await.unwrap();

            let ia2 = IaFile::open_in_path(2, file_type, &local_path, mgr.clone()).unwrap();
            let _ = ia2.multi_read_async(100, 64).await.unwrap();
            let _ = ia2.multi_read_async(100, 64).await.unwrap();

            // To make sure that segments are written to local store.
            mgr.flush_tasks(Duration::from_secs(5)).await.unwrap();
        }

        {
            let mgr = IaManager::new(options, s3fs.clone(), rt).await.unwrap();

            let mut segments_ident = mgr.get_local_segments().await;
            segments_ident.sort_by(|(m_ident, ..), (n_ident, ..)| m_ident.cmp(n_ident));
            let expected_segments = [
                // file_id, start_off, end_off
                (1, 0, 66),
                (2, 66, 132),
                (2, 132, 198),
            ]
            .iter()
            .map(|&(file_id, start_off, end_off)| FileSegmentIdent {
                file_id,
                start_off,
                end_off,
            })
            .collect::<Vec<_>>();
            assert_eq!(
                segments_ident.len(),
                expected_segments.len(),
                "{:?}",
                segments_ident
            );
            for ((ident, segment, _), expected) in segments_ident
                .into_iter()
                .zip(expected_segments.into_iter())
            {
                assert_eq!(ident, expected);
                assert_matches!(segment, FileSegmentData::InStore);
            }

            // Open file to verify meta exist.
            for file_id in 1..10 {
                let _ = IaFile::open_in_path(file_id, file_type, &local_path, mgr.clone()).unwrap();
            }
        }
    });

    oss.shutdown();
}

#[test]
fn test_abnormal_local_file() {
    init_log_for_test();

    let (temp_dir, mut oss, dfs_conf) = prepare_dfs("test");
    let temp_dir = temp_dir.path();
    let local_path = temp_dir.join("ia");

    let s3fs = S3Fs::new_from_config(dfs_conf);
    let _s3fs = s3fs.clone();

    let runtime = tokio::runtime::Runtime::new().unwrap();
    let rt = runtime.handle().clone();
    runtime.block_on(async move {
        let file_id = 42;
        let file_type = FileType::Sst;
        let (file_data, user_data, table_meta_off) =
            make_sstable(file_id, BLOCK_SIZE, 10, 7, 5, false);
        s3fs.put_object(
            s3fs.file_key(file_id, file_type),
            file_data,
            format!("{}.{}", file_id, file_type.suffix()),
        )
        .await
        .unwrap();

        let ia_cap = IaCapacity::MemoryAndDiskCap(0, local_path.clone(), 100000);
        let options = IaManagerOptionsBuilder::default()
            .capacity(ia_cap)
            .segment_size(SEGMENT_SIZE)
            .freq_update_interval(FREQ_UPDATE_INTERVAL)
            .build()
            .unwrap();
        let mgr = IaManager::new(options, s3fs.clone(), rt).await.unwrap();

        {
            let table_meta_data =
                IaFile::prepare_table_meta(file_id, file_type, table_meta_off, &local_path, &s3fs)
                    .await
                    .unwrap();
            let table_meta_file = InMemFile::new(file_id, table_meta_data);
            let ia_file =
                IaFile::open(file_id, file_type, Arc::new(table_meta_file), mgr.clone()).unwrap();
            let seg = ia_file.multi_read_async(0, user_data.len()).await.unwrap();
            assert_eq!(seg, user_data);
        }

        // Remove the local meta & segment.
        {
            tokio::fs::remove_file(table_meta_file_local_path(file_id, file_type, &local_path))
                .await
                .unwrap();

            let main_store = LocalFileStore::new(local_path.clone());
            main_store
                .remove(
                    file_id,
                    &FileSegmentIdent {
                        file_id,
                        start_off: 10,
                        end_off: 20,
                    }
                    .local_filename(),
                )
                .unwrap();
        }

        {
            // First open file failed due to local meta not found.
            IaFile::open_in_path(file_id, file_type, &local_path, mgr.clone()).unwrap_err();
            // Prepare again.
            IaFile::prepare_table_meta(file_id, file_type, table_meta_off, &local_path, &s3fs)
                .await
                .unwrap();
            let ia_file =
                IaFile::open_in_path(file_id, file_type, &local_path, mgr.clone()).unwrap();

            // Read can handle local segment not found by retry to get from remote.
            let seg = ia_file.multi_read_async(0, user_data.len()).await.unwrap();
            assert_eq!(seg, user_data);
        }
    });

    oss.shutdown();
}

fn make_sstable(
    file_id: u64,
    block_size: usize,
    n: usize,
    key_len: usize,
    val_len: usize,
    multi_ver: bool,
) -> (
    Bytes, // file_data
    Bytes, // user_data
    u64,   // meta_off
) {
    let mut rng = thread_rng();

    let mut builder = sstable::Builder::new(
        file_id,
        block_size,
        NO_COMPRESSION,
        0,
        ChecksumType::default(),
        None,
        None,
    );
    let mut val = vec![0; val_len];
    let mut ver = n as u64;
    let mut i = 0;
    for _ in 0..n {
        let key = format!("{:0key_len$}", i).into_bytes();
        rng.fill_bytes(val.as_mut_slice());
        let value_buf = Value::encode_buf(0u8, &[0], ver, &val);
        let value = Value::decode(&value_buf);
        builder.add(InnerKey::from_inner_buf(&key), &value, None);

        if multi_ver && rng.gen_ratio(1, 4) {
            ver -= 1;
        } else {
            i += 1;
            ver = n as u64;
        }
    }

    let mut buf = Vec::with_capacity(builder.estimated_size());
    let res = builder.finish(0, &mut buf);
    let file_data = Bytes::from(buf);
    let user_data = file_data.slice(0..res.meta_offset as usize);
    (file_data, user_data, res.meta_offset as u64)
}