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

use std::sync::{
    Arc,
    atomic::{AtomicU64, Ordering},
};

use async_trait::async_trait;
use cloud_encryption::MasterKey;
use file_system::IoRateLimiter;
use kvengine::{
    Engine, LOCK_CF, MetaChangeListener, MetaIterator, Options, RecoverHandler, WriteBatch,
    dfs::InMemFs, limiter::StoreLimiter,
};
use kvenginepb::{ChangeSet as PbChangeSet, Properties, Snapshot};
use kvproto::kvrpcpb::Op;
use rfstore::store::RegionSnapshot;
use security::SecurityManager;
use tempfile::TempDir;
use tikv::storage::{
    mvcc::CloudReader,
    txn::{CloudDeltaScanner, CloudStore, Store, TxnEntryScanner},
};
use tikv_kv::Statistics;
use tikv_util::Either;
use txn_types::{Key, Lock, LockType, SharedLocks, TimeStamp, TsSet};

struct EmptyMetaIter {
    engine_id: u64,
}

impl MetaIterator for EmptyMetaIter {
    fn iterate<F>(&mut self, _f: F) -> kvengine::Result<()>
    where
        F: FnMut(PbChangeSet),
    {
        Ok(())
    }

    fn engine_id(&self) -> u64 {
        self.engine_id
    }
}

#[derive(Clone)]
struct NoopRecoverer;

impl RecoverHandler for NoopRecoverer {
    fn recover(
        &self,
        _engine: &Engine,
        _shard: &Arc<kvengine::Shard>,
        _info: &kvengine::ShardMeta,
        _is_parent: bool,
    ) -> kvengine::Result<()> {
        Ok(())
    }
}

#[derive(Clone)]
struct NoopMetaChangeListener;

impl MetaChangeListener for NoopMetaChangeListener {
    fn on_change_set(&self, _cs: PbChangeSet) {}
}

#[derive(Default)]
struct TestIdAllocator {
    next_id: AtomicU64,
}

#[async_trait]
impl kvengine::IdAllocator for TestIdAllocator {
    fn alloc_id(&self, count: usize) -> kvengine::Result<Vec<u64>> {
        let start_id = self.next_id.fetch_add(count as u64, Ordering::Relaxed) + 1;
        Ok((start_id..start_id + count as u64).collect())
    }

    async fn alloc_id_async(&self, count: usize) -> kvengine::Result<Vec<u64>> {
        self.alloc_id(count)
    }
}

fn open_test_engine() -> (Engine, TempDir) {
    std::fs::create_dir_all("target").unwrap();
    let tmp_dir = tempfile::Builder::new()
        .prefix("shared-lock-scan")
        .tempdir_in("target")
        .unwrap();

    let mut opts = Options::default();
    opts.local_dirs = vec![tmp_dir.path().to_path_buf()];

    let fs = Arc::new(InMemFs::new());
    let mut meta_iter = EmptyMetaIter { engine_id: 1 };
    let rate_limiter = Arc::new(IoRateLimiter::new_for_test());
    let store_limiter = Arc::new(StoreLimiter::dummy());

    let engine = Engine::open(
        fs,
        Arc::new(opts),
        kvengine::KvEngineConfig::default(),
        &mut meta_iter,
        NoopRecoverer,
        Arc::new(TestIdAllocator::default()),
        Box::new(NoopMetaChangeListener),
        rate_limiter,
        store_limiter,
        None,
        MasterKey::new(&[1u8; 32]),
        Arc::new(SecurityManager::default()),
        true,
    )
    .unwrap();

    (engine, tmp_dir)
}

fn make_single_shard_changeset(shard_id: u64, sequence: u64) -> PbChangeSet {
    let mut props = Properties::default();
    props.set_shard_id(shard_id);

    let mut snap = Snapshot::default();
    snap.set_properties(props);
    snap.set_outer_start(Vec::new());
    snap.set_outer_end(vec![0xFF]);
    snap.set_base_version(0);
    snap.set_data_sequence(sequence);
    snap.set_inner_key_off(0);
    snap.set_max_ts(0);

    let mut cs = PbChangeSet::default();
    cs.set_shard_id(shard_id);
    cs.set_shard_ver(1);
    cs.set_sequence(sequence);
    cs.set_snapshot(snap);
    cs
}

fn write_lock_values(engine: &Engine, shard_id: u64, entries: &[(Vec<u8>, Vec<u8>)]) {
    let shard = engine.get_shard(shard_id).unwrap();
    let mut wb = WriteBatch::default();
    wb.reset(shard_id);
    wb.set_sequence(shard.get_write_sequence() + 1);
    for (raw_key, lock_value) in entries {
        wb.put(LOCK_CF, raw_key, lock_value, 0, &[], 0);
    }
    engine.write(&mut wb, &[]);
}

fn make_shared_holder(start_ts: u64) -> Lock {
    Lock::new(
        LockType::Lock,
        b"primary".to_vec(),
        start_ts.into(),
        100,
        None,
        TimeStamp::zero(),
        1,
        TimeStamp::zero(),
    )
}

fn make_shared_locks(start_tss: &[u64]) -> SharedLocks {
    let mut shared_locks = SharedLocks::new();
    for start_ts in start_tss {
        shared_locks
            .insert_lock(make_shared_holder(*start_ts))
            .unwrap();
    }
    shared_locks
}

#[test]
fn test_cloud_reader_scan_lock_reports_shared_lock_infos() {
    let (engine, _tmp_dir) = open_test_engine();
    let shard_id = 1;
    engine
        .ingest(
            kvengine::ChangeSet::new(make_single_shard_changeset(shard_id, 1)),
            true,
        )
        .unwrap();

    let raw_key = b"k_shared_scan".to_vec();
    let shared_locks = make_shared_locks(&[5, 7, 9]);
    write_lock_values(
        &engine,
        shard_id,
        &[(raw_key.clone(), shared_locks.to_bytes())],
    );

    let snapshot = engine.get_snap_access(shard_id).unwrap();
    let mut reader = CloudReader::new(snapshot, false);

    let loaded = reader.load_lock(&Key::from_raw(&raw_key)).unwrap().unwrap();
    match loaded {
        Either::Right(shared_locks) => assert_eq!(shared_locks.len(), 3),
        Either::Left(lock) => panic!("expected shared locks, got {:?}", lock),
    }

    let (locks, has_remain) = reader
        .scan_locks(None, None, None, |lock| lock.ts <= 10.into(), 0)
        .unwrap();
    assert!(!has_remain);
    assert_eq!(locks.len(), 1);

    let (key, lock_or_shared) = locks.into_iter().next().unwrap();
    let lock_info = match lock_or_shared {
        Either::Right(shared_locks) => shared_locks.into_lock_info(key.into_raw().unwrap()),
        Either::Left(lock) => panic!("expected shared locks, got {:?}", lock),
    };
    assert_eq!(lock_info.get_lock_type(), Op::SharedLock);
    assert_eq!(lock_info.get_key(), raw_key);
    assert_eq!(lock_info.get_shared_lock_infos().len(), 3);
}

#[test]
fn test_cloud_reader_scan_lock_limit_counts_shared_holders() {
    let (engine, _tmp_dir) = open_test_engine();
    let shard_id = 1;
    engine
        .ingest(
            kvengine::ChangeSet::new(make_single_shard_changeset(shard_id, 1)),
            true,
        )
        .unwrap();

    let raw_key = b"k_shared_limit".to_vec();
    let shared_locks = make_shared_locks(&[11, 12, 13]);
    write_lock_values(&engine, shard_id, &[(raw_key, shared_locks.to_bytes())]);

    let snapshot = engine.get_snap_access(shard_id).unwrap();
    let mut reader = CloudReader::new(snapshot, false);
    let (locks, has_remain) = reader
        .scan_locks(None, None, None, |lock| lock.ts <= 20.into(), 2)
        .unwrap();

    assert!(has_remain);
    assert_eq!(locks.len(), 1);

    let (_key, lock_or_shared) = locks.into_iter().next().unwrap();
    match lock_or_shared {
        Either::Right(shared_locks) => assert_eq!(shared_locks.len(), 2),
        Either::Left(lock) => panic!("expected shared locks, got {:?}", lock),
    }
}

#[test]
fn test_cloud_store_shared_lock_wrapper_is_read_compatible() {
    let (engine, _tmp_dir) = open_test_engine();
    let shard_id = 1;
    engine
        .ingest(
            kvengine::ChangeSet::new(make_single_shard_changeset(shard_id, 1)),
            true,
        )
        .unwrap();

    let raw_key = b"k_shared_cloud_store".to_vec();
    let shared_locks = make_shared_locks(&[5, 7]);
    write_lock_values(
        &engine,
        shard_id,
        &[(raw_key.clone(), shared_locks.to_bytes())],
    );

    let snapshot = engine.get_snap_access(shard_id).unwrap();
    let region_snapshot = RegionSnapshot::from_snapshot(snapshot, None, None);
    let store = CloudStore::new(region_snapshot, u64::MAX, TsSet::Empty, false, false);

    let mut statistics = Statistics::default();
    let value = store
        .get(&Key::from_raw(&raw_key), &mut statistics)
        .unwrap();
    assert_eq!(value, None);
    store.check_locks_in_range(&[], &[0xFF], false).unwrap();
}

#[test]
fn test_cloud_delta_scanner_skips_shared_lock_wrapper() {
    let (engine, _tmp_dir) = open_test_engine();
    let shard_id = 1;
    engine
        .ingest(
            kvengine::ChangeSet::new(make_single_shard_changeset(shard_id, 1)),
            true,
        )
        .unwrap();

    let raw_key = b"k_shared_delta".to_vec();
    let shared_locks = make_shared_locks(&[11, 12]);
    write_lock_values(&engine, shard_id, &[(raw_key, shared_locks.to_bytes())]);

    let snapshot = engine.get_snap_access(shard_id).unwrap();
    let mut scanner = CloudDeltaScanner::new(snapshot, 0, None, None, false);
    scanner.init().unwrap();
    assert!(scanner.next_entry().unwrap().is_none());
}