// Copyright 2026 PingCAP, Inc.

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

use async_trait::async_trait;
use cloud_encryption::MasterKey;
use file_system::IoRateLimiter;
use kvengine::{
    EXTRA_CF, Engine, MetaChangeListener, MetaIterator, Options, RecoverHandler, UserMeta,
    WriteBatch, dfs::InMemFs, limiter::StoreLimiter,
};
use kvenginepb::{ChangeSet as PbChangeSet, Properties, Snapshot};
use security::SecurityManager;
use tempfile::TempDir;
use tikv::storage::mvcc::{CloudReader, TxnCommitRecord};
use txn_types::{Key, TimeStamp, WriteType};

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("extra-cf-compat")
        .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
}

#[test]
fn test_extra_cf_get_extra_backward_compatible() {
    let (engine, _tmp_dir) = open_test_engine();

    let shard_id = 1;
    let start_seq = 1;
    let pb_cs = make_single_shard_changeset(shard_id, start_seq);
    engine
        .ingest(kvengine::ChangeSet::new(pb_cs), true)
        .unwrap();

    let raw_key = b"xkey_00000001".to_vec();
    let start_ts = 100;
    let commit_ts = 120;

    // Detect which encoding is currently used by
    // `kvengine::encode_extra_txn_status_key` and write the record using the
    // *other* encoding. This ensures the read path exercises the
    // backward-compatible fallback regardless of which encoding is primary.
    let key_desc = kvengine::encode_extra_txn_status_key_desc(&raw_key, start_ts);
    let key_revbits = kvengine::encode_extra_txn_status_key_revbits(&raw_key, start_ts);
    let primary_key = kvengine::encode_extra_txn_status_key(&raw_key, start_ts);
    assert_eq!(primary_key.as_ref(), key_desc.as_ref());
    let non_primary_key = if primary_key.as_ref() == key_desc.as_ref() {
        key_revbits.to_vec()
    } else {
        key_desc.to_vec()
    };
    assert_ne!(primary_key.as_ref(), non_primary_key.as_slice());

    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);
    let user_meta = UserMeta::new(start_ts, commit_ts).to_array();
    wb.put(EXTRA_CF, &non_primary_key, &[0], 0, &user_meta, commit_ts);
    engine.write(&mut wb, &[]);

    let snapshot = engine.get_snap_access(shard_id).unwrap();
    let mut reader = CloudReader::new(snapshot, false);
    let record = reader
        .get_txn_commit_record(&Key::from_raw(&raw_key), TimeStamp::new(start_ts))
        .expect("extra lookup succeeds");

    let (ts, write) = match record {
        TxnCommitRecord::SingleRecord { commit_ts, write } => (commit_ts, write),
        _ => panic!("unexpected txn commit record: {:?}", record),
    };

    assert_eq!(ts.into_inner(), commit_ts);
    assert_eq!(write.write_type, WriteType::Lock);
    assert_eq!(write.start_ts.into_inner(), start_ts);
}

/// Verify that rollback-style EXTRA_CF records (commit_ts == 0) written with
/// the legacy encoding are still readable after switching the writer to desc.
#[test]
fn test_extra_cf_get_extra_backward_compatible_rollback() {
    let (engine, _tmp_dir) = open_test_engine();

    let shard_id = 1;
    let start_seq = 1;
    let pb_cs = make_single_shard_changeset(shard_id, start_seq);
    engine
        .ingest(kvengine::ChangeSet::new(pb_cs), true)
        .unwrap();

    let raw_key = b"xkey_00000002".to_vec();
    let start_ts = 200;
    let commit_ts = 0; // rollback record

    // Write using the legacy (non-primary) encoding to simulate a historical
    // rollback record written before the encoding switch.
    let key_revbits = kvengine::encode_extra_txn_status_key_revbits(&raw_key, start_ts);
    let primary_key = kvengine::encode_extra_txn_status_key(&raw_key, start_ts);
    // Confirm primary is now desc, so revbits is the legacy path.
    assert_ne!(primary_key.as_ref(), key_revbits.as_ref());

    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);
    let user_meta = UserMeta::new(start_ts, commit_ts).to_array();
    wb.put(
        EXTRA_CF,
        key_revbits.as_ref(),
        &[0],
        0,
        &user_meta,
        start_ts, // rollback version = start_ts
    );
    engine.write(&mut wb, &[]);

    let snapshot = engine.get_snap_access(shard_id).unwrap();
    let mut reader = CloudReader::new(snapshot, false);
    let record = reader
        .get_txn_commit_record(&Key::from_raw(&raw_key), TimeStamp::new(start_ts))
        .expect("extra lookup succeeds");

    let (ts, write) = match record {
        TxnCommitRecord::SingleRecord { commit_ts, write } => (commit_ts, write),
        _ => panic!("unexpected txn commit record: {:?}", record),
    };

    // For rollback records, get_txn_commit_record returns (start_ts, Rollback).
    assert_eq!(ts.into_inner(), start_ts);
    assert_eq!(write.write_type, WriteType::Rollback);
    assert_eq!(write.start_ts.into_inner(), start_ts);
}