// Copyright 2025 PingCAP, Inc.

use std::{rc::Rc, sync::Arc, time::Duration};

use futures::executor::block_on;
use test_util::init_log_for_test;
use tikv_util::{deadline::Deadline, time::Instant};
use tracker::{
    GLOBAL_TRACKERS, RequestInfo, RequestType, clear_tls_tracker_token, set_tls_tracker_token,
};

use crate::{
    Iterator,
    ia::{ia_file::IaFile, manager::IaManager, util::IaManagerOptionsBuilder},
    next, next_async,
    table::{
        file::{File, InMemFile},
        sstable::{BlockCache, Footer, NewSsTableCtx, SsTable},
        tiny_meta::SstTinyMeta,
    },
    tests::{new_table, new_test_engine_opt},
    util::test_util::KeyBuilder,
};

#[rstest::rstest]
#[case::concurrency(2)]
#[case::concurrency(8)]
fn test_sync_read(#[case] concurrency: usize) {
    init_log_for_test();

    let file_id = 1;
    let (engine, _) = new_test_engine_opt(true, 1024, "");
    let mut saved_vals: Vec<Rc<Vec<u8>>> = Vec::new();
    let t = new_table(&engine, file_id, 0, 100, 1000, false, &mut saved_vals);

    let ia_rt = tokio::runtime::Builder::new_multi_thread()
        .thread_name("ia-mgr")
        .enable_all()
        .worker_threads(2)
        .build()
        .unwrap();
    let sync_read_concurrency = concurrency.div_ceil(2);
    let options = IaManagerOptionsBuilder::default()
        .segment_size(4096)
        .sync_read(false, sync_read_concurrency, Duration::from_secs(3))
        .build()
        .unwrap();
    let mgr = IaManager::new(options, engine.fs.clone(), None, ia_rt.into()).unwrap();
    let t_async = convert_local_sst_to_ia(&t, mgr.clone());

    // Test in tokio async context.
    // Simulate unified pool.
    let unified_pool = tokio::runtime::Builder::new_multi_thread()
        .thread_name("sync-read-unified-pool")
        .enable_all()
        .worker_threads(concurrency)
        .build()
        .unwrap();
    let start_time = Instant::now_coarse();
    let mut js = tokio::task::JoinSet::new();
    for i in 0..concurrency {
        let kb = engine.key_builder.clone();
        let tt = if i & 1 == 0 {
            t_async.clone()
        } else {
            t.clone()
        };
        let task = unified_pool.spawn(async move {
            verify_table_by_scan(&kb, &tt, 0, 100, 1000);
        });
        js.spawn_on(task, unified_pool.handle());
    }
    block_on(async move {
        while let Some(r) = js.join_next().await {
            r.unwrap().unwrap();
        }
    });
    println!(
        "test_sync_read: concurrency {}, take {:?}",
        concurrency,
        start_time.saturating_elapsed()
    );

    // Test in sync context.
    verify_table_by_scan(&engine.key_builder, &t_async, 0, 100, 1000);
}

#[maybe_async::both]
#[track_caller]
pub(crate) async fn verify_table_by_scan(
    kb: &KeyBuilder,
    t: &SsTable,
    begin: usize,
    end: usize,
    ver: u64,
) {
    let mut i = begin;
    let mut it = t.new_iterator(false, false);
    it.rewind().await;
    while it.valid() {
        let key = it.key();
        assert_eq!(kb.i_to_inner_key(i).as_ref(), key);
        let val = it.value();
        assert_eq!(val.version, ver);
        assert_eq!(val.get_value(), key.repeat(2));

        next!(it).await;
        i += 1;
    }
    assert_eq!(i, end);
}

pub(crate) fn open_ia_file_from_sst(t: &SsTable, mgr: IaManager) -> IaFile {
    let meta_data = t
        .file()
        .read(
            t.meta_offset() as u64,
            t.size() as usize - t.meta_offset() as usize,
        )
        .unwrap();
    let meta_file = InMemFile::new(t.id(), meta_data);
    IaFile::open_for_sst(t.id(), Arc::new(meta_file), mgr, None, true).unwrap()
}

pub(crate) fn convert_local_sst_to_ia(t: &SsTable, mgr: IaManager) -> SsTable {
    let ia_file = open_ia_file_from_sst(t, mgr);
    SsTable::new(
        Arc::new(ia_file),
        BlockCache::None,
        cloud_encryption::EncryptionState::Disabled,
    )
    .unwrap()
}

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

    let file_id = 1;
    let (engine, _) = new_test_engine_opt(true, 256, "");
    let mut saved_vals: Vec<Rc<Vec<u8>>> = Vec::new();
    let t = new_table(&engine, file_id, 0, 100, 1000, false, &mut saved_vals);

    let ia_rt = tokio::runtime::Builder::new_multi_thread()
        .thread_name("ia-mgr")
        .enable_all()
        .worker_threads(1)
        .build()
        .unwrap();
    let options = IaManagerOptionsBuilder::default()
        .segment_size(512)
        .build()
        .unwrap();
    let mgr = IaManager::new(options, engine.fs.clone(), None, ia_rt.into()).unwrap();

    let ia_file = Arc::new(open_ia_file_from_sst(&t, mgr.clone()));
    let meta_file = ia_file.table_meta_file().clone();

    let mut ctx = NewSsTableCtx::default();
    let ia_table = SsTable::new_with_ctx(
        ia_file.clone(),
        BlockCache::None,
        cloud_encryption::EncryptionState::Disabled,
        &mut ctx,
    )
    .unwrap();

    let (footer, props_data) = ctx
        .footer_and_properties_data
        .expect("footer and props data");
    let sst_tiny_meta = SstTinyMeta::from_sstable(&ia_table, footer, props_data.as_ref());
    let segment_offsets = sst_tiny_meta
        .segment_offsets
        .as_ref()
        .expect("segment offsets should be present");
    assert!(!segment_offsets.is_empty());
    let ia_with_tiny =
        IaFile::open_for_sst(t.id(), meta_file, mgr, Some(&sst_tiny_meta), true).unwrap();

    assert_eq!(ia_file.id, ia_with_tiny.id);
    assert_eq!(ia_file.size, ia_with_tiny.size);
    assert_eq!(ia_file.ftype, ia_with_tiny.ftype);
    assert_eq!(ia_file.table_meta_off, ia_with_tiny.table_meta_off);
    let segment_offsets_u64: Vec<u64> = segment_offsets.iter().copied().map(u64::from).collect();
    assert_eq!(ia_file.segment_offsets, segment_offsets_u64);
    assert_eq!(ia_file.segment_offsets, ia_with_tiny.segment_offsets);
}

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

    let file_id = 1;
    let (engine, _) = new_test_engine_opt(true, 256, "");
    let mut saved_vals: Vec<Rc<Vec<u8>>> = Vec::new();
    let table = new_table(&engine, file_id, 0, 100, 1000, false, &mut saved_vals);
    let manager = engine.new_ia_manager(IaManagerOptionsBuilder::default().build().unwrap());

    let mut meta_data = table
        .file()
        .read(
            table.meta_offset() as u64,
            table.size() as usize - table.meta_offset() as usize,
        )
        .unwrap()
        .to_vec();
    let footer_offset = meta_data.len() - SsTable::footer_size();
    let mut footer = Footer::default();
    footer.unmarshal(&meta_data[footer_offset..]);
    let table_size = footer.meta_offset() as usize + meta_data.len();
    // Keep the offsets ordered but move properties one byte into the footer.
    footer.properties_offset = (table_size - SsTable::footer_size() + 1) as u32;
    let mut footer_data = Vec::with_capacity(SsTable::footer_size());
    footer.marshal(&mut footer_data);
    meta_data[footer_offset..].copy_from_slice(&footer_data);

    let meta_file = Arc::new(InMemFile::new(file_id, meta_data.into()));
    assert!(matches!(
        IaFile::open_for_sst(file_id, meta_file, manager, None, true),
        Err(crate::table::Error::InvalidFileSize)
    ));
}

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

    let file_id = 1;
    let (engine, _) = new_test_engine_opt(true, 256, "");
    let mut saved_vals: Vec<Rc<Vec<u8>>> = Vec::new();
    let t = new_table(&engine, file_id, 0, 8, 1000, false, &mut saved_vals);

    let ia_rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let options = IaManagerOptionsBuilder::default()
        .segment_size(4096)
        .build()
        .unwrap();
    let mgr = IaManager::new(options, engine.fs.clone(), None, ia_rt.into()).unwrap();
    let ia_file = open_ia_file_from_sst(&t, mgr);
    let expected_segment_bytes = ia_file.get_segment_ident(0).unwrap().size();

    let token = GLOBAL_TRACKERS.insert(tracker::Tracker::new(RequestInfo {
        region_id: 1,
        start_ts: 0,
        task_id: 1,
        resource_group_tag: vec![],
        request_type: RequestType::CoprocessorDag,
        is_remote: false,
    }));
    set_tls_tracker_token(token);
    let first_read = ia_file.read_async(0, 16);
    let first_data = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(first_read)
        .unwrap();
    assert_eq!(first_data.len(), 16);
    let first_metrics = GLOBAL_TRACKERS
        .with_tracker(token, |tracker| tracker.metrics.clone())
        .unwrap();
    clear_tls_tracker_token();
    GLOBAL_TRACKERS.remove(token);

    assert_eq!(first_metrics.ia_remote_read_segment_count, 1);
    assert_eq!(
        first_metrics.ia_remote_read_segment_bytes,
        expected_segment_bytes
    );
    assert_eq!(first_metrics.ia_cache_hit_count, 0);

    let token = GLOBAL_TRACKERS.insert(tracker::Tracker::new(RequestInfo {
        region_id: 1,
        start_ts: 0,
        task_id: 2,
        resource_group_tag: vec![],
        request_type: RequestType::CoprocessorDag,
        is_remote: false,
    }));
    set_tls_tracker_token(token);
    let second_read = ia_file.read_async(0, 16);
    let second_data = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(second_read)
        .unwrap();
    assert_eq!(second_data.len(), 16);
    let second_metrics = GLOBAL_TRACKERS
        .with_tracker(token, |tracker| tracker.metrics.clone())
        .unwrap();
    clear_tls_tracker_token();
    GLOBAL_TRACKERS.remove(token);

    assert_eq!(second_metrics.ia_remote_read_segment_count, 0);
    assert_eq!(second_metrics.ia_remote_read_segment_bytes, 0);
    assert_eq!(second_metrics.ia_remote_read_segment_nanos, 0);
    assert_eq!(second_metrics.ia_cache_hit_count, 1);
}

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

    let file_id = 1;
    let (engine, _) = new_test_engine_opt(true, 256, "");
    let mut saved_vals: Vec<Rc<Vec<u8>>> = Vec::new();
    let t = new_table(&engine, file_id, 0, 8, 1000, false, &mut saved_vals);

    let ia_rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let options = IaManagerOptionsBuilder::default()
        .segment_size(4096)
        .build()
        .unwrap();
    let mgr = IaManager::new(options, engine.fs.clone(), None, ia_rt.into()).unwrap();
    let ia_file = open_ia_file_from_sst(&t, mgr);
    let expected_segment_bytes = ia_file.get_segment_ident(0).unwrap().size();

    let token = GLOBAL_TRACKERS.insert(tracker::Tracker::new(RequestInfo {
        region_id: 1,
        start_ts: 0,
        task_id: 1,
        resource_group_tag: vec![],
        request_type: RequestType::CoprocessorDag,
        is_remote: false,
    }));
    set_tls_tracker_token(token);
    let first_read = ia_file.mmap_range(0, 16);
    let (first_data, _) = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(first_read)
        .unwrap();
    assert_eq!(first_data.len(), 16);
    let first_metrics = GLOBAL_TRACKERS
        .with_tracker(token, |tracker| tracker.metrics.clone())
        .unwrap();
    clear_tls_tracker_token();
    GLOBAL_TRACKERS.remove(token);

    assert_eq!(first_metrics.ia_remote_read_segment_count, 1);
    assert_eq!(
        first_metrics.ia_remote_read_segment_bytes,
        expected_segment_bytes
    );
    assert_eq!(first_metrics.ia_cache_hit_count, 0);

    let token = GLOBAL_TRACKERS.insert(tracker::Tracker::new(RequestInfo {
        region_id: 1,
        start_ts: 0,
        task_id: 2,
        resource_group_tag: vec![],
        request_type: RequestType::CoprocessorDag,
        is_remote: false,
    }));
    set_tls_tracker_token(token);
    let second_read = ia_file.mmap_range(0, 16);
    let (second_data, _) = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(second_read)
        .unwrap();
    assert_eq!(second_data.len(), 16);
    let second_metrics = GLOBAL_TRACKERS
        .with_tracker(token, |tracker| tracker.metrics.clone())
        .unwrap();
    clear_tls_tracker_token();
    GLOBAL_TRACKERS.remove(token);

    assert_eq!(second_metrics.ia_remote_read_segment_count, 0);
    assert_eq!(second_metrics.ia_remote_read_segment_bytes, 0);
    assert_eq!(second_metrics.ia_remote_read_segment_nanos, 0);
    assert_eq!(second_metrics.ia_cache_hit_count, 1);
}

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

    let file_id = 1;
    let (engine, _) = new_test_engine_opt(true, 256, "");
    let mut saved_vals: Vec<Rc<Vec<u8>>> = Vec::new();
    let t = new_table(&engine, file_id, 0, 8, 1000, false, &mut saved_vals);

    let ia_rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();
    let options = IaManagerOptionsBuilder::default()
        .segment_size(4096)
        .build()
        .unwrap();
    let mgr = IaManager::new(options, engine.fs.clone(), None, ia_rt.into()).unwrap();
    let ia_file = open_ia_file_from_sst(&t, mgr);
    let ident = ia_file.get_segment_ident(0).unwrap();
    let expected_segment_bytes = ident.size();

    let token = GLOBAL_TRACKERS.insert(tracker::Tracker::new(RequestInfo {
        region_id: 1,
        start_ts: 0,
        task_id: 1,
        resource_group_tag: vec![],
        request_type: RequestType::CoprocessorDag,
        is_remote: false,
    }));
    set_tls_tracker_token(token);
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(async {
            ia_file
                .mgr
                .prefetch_segment(
                    ident,
                    ia_file.ftype,
                    0,
                    Deadline::from_now(Duration::from_secs(3)),
                    None,
                    true,
                )
                .await
                .unwrap();
            let data = ia_file.read_async(0, 16).await.unwrap();
            assert_eq!(data.len(), 16);
        });
    let metrics = GLOBAL_TRACKERS
        .with_tracker(token, |tracker| tracker.metrics.clone())
        .unwrap();
    clear_tls_tracker_token();
    GLOBAL_TRACKERS.remove(token);

    assert_eq!(metrics.ia_remote_read_segment_count, 1);
    assert_eq!(metrics.ia_remote_read_segment_bytes, expected_segment_bytes);
    assert_eq!(metrics.ia_cache_hit_count, 1);
}