// Copyright 2026 PingCAP, Inc.

use std::sync::Arc;

use bytes::Bytes;
use futures::{SinkExt, executor::block_on};
use kvengine::{
    dfs::{self, Dfs, FileType},
    table::{
        InnerKey, PessimisticAction, TXN_FILE_FORMAT_V2,
        txn_file::{OP_DELETE, OP_PUT, TxnChunkBuilder},
    },
};
use kvproto::{cdcpb::*, kvrpcpb::*};
use pd_client::PdClient;
use test_raftstore::sleep_ms;
use txn_types::Key;

use crate::{
    TestCluster, TestSuite, get_ready_region, new_event_feed_v2, subscribe_until_initialized,
};

const MAGIC_KEYSPACE_START: &[u8] = b"xkey";
const MAGIC_KEYSPACE_END: &[u8] = b"xkez";

fn txn_file_key(i: usize) -> Vec<u8> {
    let mut key = MAGIC_KEYSPACE_START.to_vec();
    key.extend(format!("t_key_{i:08}").as_bytes());
    key
}

fn txn_file_value(i: usize) -> Vec<u8> {
    format!("txn-file-value-{i:08}").into_bytes()
}

fn txn_file_delete_value() -> Vec<u8> {
    Vec::new()
}

fn txn_file_index_key(i: usize) -> Vec<u8> {
    let mut key = Vec::new();
    key.push(b't');
    key.extend_from_slice(&1_i64.to_be_bytes());
    key.extend_from_slice(b"_i");
    key.extend_from_slice(format!("idx_{i:08}").as_bytes());
    key
}

fn put_mutation(key: &[u8], value: &[u8]) -> Mutation {
    let mut mutation = Mutation::default();
    mutation.set_op(Op::Put);
    mutation.key = key.to_vec();
    mutation.value = value.to_vec();
    mutation
}

fn get_dfs(suite: &TestSuite) -> Arc<dyn Dfs> {
    match &suite.cluster {
        TestCluster::CloudEngine { server, .. } => server.get_dfs().unwrap(),
        TestCluster::Tikv(_) => panic!("txn-file CDC tests require cloud-engine"),
    }
}

fn build_txn_file_chunk(
    dfs: &Arc<dyn Dfs>,
    chunk_id: u64,
    rows: &[(Vec<u8>, Vec<u8>)],
) -> Vec<u64> {
    let rows = rows
        .iter()
        .map(|(key, value)| (key.clone(), OP_PUT, value.clone()))
        .collect::<Vec<_>>();
    build_txn_file_chunk_with_ops(dfs, chunk_id, &rows)
}

fn build_txn_file_chunk_with_ops(
    dfs: &Arc<dyn Dfs>,
    chunk_id: u64,
    rows: &[(Vec<u8>, u8, Vec<u8>)],
) -> Vec<u64> {
    let mut builder = TxnChunkBuilder::new(chunk_id, 64, None, TXN_FILE_FORMAT_V2);
    for (key, op, value) in rows {
        builder.add_entry(
            InnerKey::from_outer_key(key),
            *op,
            PessimisticAction::SkipCheck,
            value,
        );
    }
    let mut data = Vec::new();
    builder.finish(&mut data);
    let opts = dfs::Options::default().with_type(FileType::TxnChunk);
    dfs.get_runtime()
        .block_on(dfs.create(chunk_id, Bytes::from(data), opts))
        .unwrap();
    vec![chunk_id]
}

fn build_txn_file_chunks_with_ops(
    dfs: &Arc<dyn Dfs>,
    first_chunk_id: u64,
    chunks: &[Vec<(Vec<u8>, u8, Vec<u8>)>],
) -> Vec<u64> {
    let mut chunk_ids = Vec::with_capacity(chunks.len());
    for (offset, rows) in chunks.iter().enumerate() {
        chunk_ids.extend(build_txn_file_chunk_with_ops(
            dfs,
            first_chunk_id + offset as u64,
            rows,
        ));
    }
    chunk_ids
}

fn must_txn_file_prewrite(
    suite: &mut TestSuite,
    region_id: u64,
    primary_key: Vec<u8>,
    start_ts: u64,
    chunk_ids: Vec<u64>,
) {
    let mut req = PrewriteRequest::new();
    req.set_context(suite.get_context(region_id));
    req.set_txn_file_chunks(chunk_ids);
    req.set_primary_lock(primary_key);
    req.set_lock_ttl(6000);
    req.set_start_version(start_ts);
    req.set_min_commit_ts(start_ts + 1);

    let resp = suite.get_tikv_client(region_id).kv_prewrite(&req).unwrap();
    assert!(!resp.has_region_error(), "{:?}", resp.get_region_error());
    assert!(resp.get_errors().is_empty(), "{:?}", resp.get_errors());
}

fn must_txn_file_commit(suite: &mut TestSuite, region_id: u64, start_ts: u64, commit_ts: u64) {
    let mut req = CommitRequest::new();
    req.set_context(suite.get_context(region_id));
    req.set_start_version(start_ts);
    req.set_is_txn_file(true);
    req.set_commit_version(commit_ts);

    let resp = suite.get_tikv_client(region_id).kv_commit(&req).unwrap();
    assert!(!resp.has_region_error(), "{:?}", resp.get_region_error());
    assert!(!resp.has_error(), "{:?}", resp.get_error());
}

fn must_txn_file_rollback(suite: &mut TestSuite, region_id: u64, start_ts: u64) {
    let mut req = BatchRollbackRequest::new();
    req.set_context(suite.get_context(region_id));
    req.set_is_txn_file(true);
    req.set_start_version(start_ts);

    let resp = suite
        .get_tikv_client(region_id)
        .kv_batch_rollback(&req)
        .unwrap();
    assert!(!resp.has_region_error(), "{:?}", resp.get_region_error());
    assert!(!resp.has_error(), "{:?}", resp.get_error());
}

fn collect_rows(
    receive_event: &(dyn Fn(bool) -> ChangeDataEvent + Send),
    expected_type: EventLogType,
    expected_len: usize,
) -> Vec<EventRow> {
    let mut rows = Vec::new();
    for _ in 0..10 {
        let event = receive_event(true);
        if event.has_resolved_ts() {
            continue;
        }
        for event in event.events.into_iter() {
            match event.event.unwrap() {
                Event_oneof_event::Entries(mut entries) => {
                    for row in entries.take_entries().into_vec() {
                        assert_eq!(row.get_type(), expected_type, "{:?}", row);
                        rows.push(row);
                    }
                }
                Event_oneof_event::Error(err) => panic!("unexpected CDC error: {:?}", err),
                other => panic!("unexpected CDC event: {:?}", other),
            }
        }
        if rows.len() == expected_len {
            return rows;
        }
        assert!(rows.len() < expected_len, "too many rows: {:?}", rows);
    }
    panic!(
        "timed out waiting for {expected_len} {:?} rows, got {:?}",
        expected_type, rows
    );
}

fn collect_entry_batches(
    receive_event: &(dyn Fn(bool) -> ChangeDataEvent + Send),
    expected_type: EventLogType,
    expected_len: usize,
) -> Vec<Vec<EventRow>> {
    let mut batches = Vec::new();
    let mut row_count = 0;
    for _ in 0..20 {
        let event = receive_event(true);
        if event.has_resolved_ts() {
            continue;
        }
        for event in event.events.into_iter() {
            match event.event.unwrap() {
                Event_oneof_event::Entries(mut entries) => {
                    let rows = entries.take_entries().into_vec();
                    for row in &rows {
                        assert_eq!(row.get_type(), expected_type, "{:?}", row);
                    }
                    row_count += rows.len();
                    batches.push(rows);
                }
                Event_oneof_event::Error(err) => panic!("unexpected CDC error: {:?}", err),
                other => panic!("unexpected CDC event: {:?}", other),
            }
        }
        if row_count == expected_len {
            return batches;
        }
        assert!(row_count < expected_len, "too many rows: {:?}", batches);
    }
    panic!(
        "timed out waiting for {expected_len} {:?} rows, got {:?}",
        expected_type, batches
    );
}

fn collect_rows_until_epoch_not_match(
    receive_event: &(dyn Fn(bool) -> ChangeDataEvent + Send),
    expected_type: EventLogType,
    expected_len: usize,
) -> (Vec<EventRow>, Vec<kvproto::metapb::Region>) {
    let mut rows = Vec::new();
    for _ in 0..20 {
        let event = receive_event(true);
        if event.has_resolved_ts() {
            continue;
        }
        for event in event.events.into_iter() {
            match event.event.unwrap() {
                Event_oneof_event::Entries(mut entries) => {
                    for row in entries.take_entries().into_vec() {
                        assert_eq!(row.get_type(), expected_type, "{:?}", row);
                        rows.push(row);
                    }
                    assert!(rows.len() <= expected_len, "too many rows: {:?}", rows);
                }
                Event_oneof_event::Error(err) => {
                    assert!(err.has_epoch_not_match(), "{:?}", err);
                    assert_eq!(rows.len(), expected_len, "{:?}", rows);
                    return (
                        rows,
                        err.get_epoch_not_match().get_current_regions().to_vec(),
                    );
                }
                other => panic!("unexpected CDC event: {:?}", other),
            }
        }
    }
    panic!(
        "timed out waiting for epoch-not-match after rows {:?}",
        rows
    );
}

fn assert_resolved_after(
    receive_event: &(dyn Fn(bool) -> ChangeDataEvent + Send),
    region_id: u64,
    min_ts: u64,
) {
    let mut previous_ts = 0;
    for _ in 0..30 {
        let event = receive_event(true);
        assert!(
            event.events.is_empty(),
            "unexpected row events: {:?}",
            event
        );
        if event.has_resolved_ts() {
            let resolved_ts = event.get_resolved_ts();
            assert!(resolved_ts.ts >= previous_ts, "{:?}", resolved_ts);
            assert_eq!(resolved_ts.regions, vec![region_id], "{:?}", resolved_ts);
            previous_ts = resolved_ts.ts;
            if resolved_ts.ts > min_ts {
                return;
            }
        }
    }
    panic!("resolved ts did not advance past {min_ts}");
}

#[test]
fn test_txn_file_cdc_expands_commit_rows_in_order_and_advances_resolved_ts() {
    let mut suite = TestSuite::new(1, ApiVersion::V2);
    let key1 = txn_file_key(1);
    let key2 = txn_file_key(2);
    let region = get_ready_region(&mut suite, &key1);
    let region_id = region.id;

    let init_start_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let init_commit_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    suite.must_kv_prewrite(
        region_id,
        vec![put_mutation(&key1, b"old-value")],
        key1.clone(),
        init_start_ts,
    );
    suite.must_kv_commit(region_id, vec![key1.clone()], init_start_ts, init_commit_ts);

    let checkpoint_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let (mut req_tx, _event_feed_wrap, receive_event) =
        new_event_feed_v2(suite.get_region_cdc_client(region_id));
    let mut req = suite.new_changedata_request(region_id);
    req.request_id = 1;
    req.checkpoint_ts = checkpoint_ts.into_inner();
    req.set_extra_op(ExtraOp::ReadOldValue);
    req.set_start_key(Key::from_raw(MAGIC_KEYSPACE_START).into_encoded());
    req.set_end_key(Key::from_raw(MAGIC_KEYSPACE_END).into_encoded());
    subscribe_until_initialized(
        &mut req_tx,
        receive_event.as_ref(),
        &req,
        "test_txn_file_cdc_expands_commit_rows_in_order_and_advances_resolved_ts initial subscribe",
    );

    let start_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let rows = vec![
        (key1.clone(), b"updated-value".to_vec()),
        (key2.clone(), txn_file_value(2)),
    ];
    let chunk_ids = build_txn_file_chunk(&get_dfs(&suite), start_ts.into_inner() + 1, &rows);
    must_txn_file_prewrite(
        &mut suite,
        region_id,
        key1.clone(),
        start_ts.into_inner(),
        chunk_ids,
    );

    let commit_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    must_txn_file_commit(
        &mut suite,
        region_id,
        start_ts.into_inner(),
        commit_ts.into_inner(),
    );

    let commit_rows = collect_rows(receive_event.as_ref(), EventLogType::Committed, 2);
    assert_eq!(commit_rows[0].get_key(), key1.as_slice());
    assert_eq!(commit_rows[0].get_value(), b"updated-value");
    assert_eq!(commit_rows[0].get_old_value(), b"old-value");
    assert_eq!(commit_rows[1].get_key(), key2.as_slice());
    assert_eq!(commit_rows[1].get_value(), txn_file_value(2).as_slice());
    assert_eq!(commit_rows[1].get_old_value(), b"");
    for row in &commit_rows {
        assert_eq!(row.get_start_ts(), start_ts.into_inner(), "{:?}", row);
        assert_eq!(row.get_commit_ts(), commit_ts.into_inner(), "{:?}", row);
        assert_eq!(row.get_op_type(), EventRowOpType::Put, "{:?}", row);
    }

    assert_resolved_after(receive_event.as_ref(), region_id, commit_ts.into_inner());
    block_on(req_tx.close()).unwrap();
    suite.stop();
}

#[test]
fn test_txn_file_cdc_rollback_emits_primary_only_and_advances_resolved_ts() {
    let mut suite = TestSuite::new(1, ApiVersion::V2);
    let key1 = txn_file_key(11);
    let key2 = txn_file_key(12);
    let region = get_ready_region(&mut suite, &key1);
    let region_id = region.id;

    let checkpoint_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let (mut req_tx, _event_feed_wrap, receive_event) =
        new_event_feed_v2(suite.get_region_cdc_client(region_id));
    let mut req = suite.new_changedata_request(region_id);
    req.request_id = 2;
    req.checkpoint_ts = checkpoint_ts.into_inner();
    req.set_start_key(Key::from_raw(MAGIC_KEYSPACE_START).into_encoded());
    req.set_end_key(Key::from_raw(MAGIC_KEYSPACE_END).into_encoded());
    subscribe_until_initialized(
        &mut req_tx,
        receive_event.as_ref(),
        &req,
        "test_txn_file_cdc_rollback_emits_primary_only_and_advances_resolved_ts initial subscribe",
    );

    let start_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let rows = vec![
        (key1.clone(), txn_file_value(11)),
        (key2.clone(), txn_file_value(12)),
    ];
    let chunk_ids = build_txn_file_chunk(&get_dfs(&suite), start_ts.into_inner() + 1, &rows);
    must_txn_file_prewrite(
        &mut suite,
        region_id,
        key1.clone(),
        start_ts.into_inner(),
        chunk_ids,
    );

    must_txn_file_rollback(&mut suite, region_id, start_ts.into_inner());
    // The normal rollback request emits the primary rollback event, while
    // txn-file rollback expansion only untracks the chunk locks internally and
    // does not emit the full chunk row set.
    let rollback_rows = collect_rows(receive_event.as_ref(), EventLogType::Rollback, 1);
    assert_eq!(rollback_rows[0].get_key(), key1.as_slice());
    assert_eq!(rollback_rows[0].get_start_ts(), start_ts.into_inner());

    // Resolved-ts advancement proves the txn-file locks were removed.
    assert_resolved_after(receive_event.as_ref(), region_id, start_ts.into_inner());

    block_on(req_tx.close()).unwrap();
    // Give the CDC service a short chance to process stream shutdown before the
    // test cluster tears down its gRPC environment.
    sleep_ms(20);
    suite.stop();
}

#[test]
fn test_txn_file_cdc_multi_chunk_split_batches_delete_and_filter_index_rows() {
    let mut suite = TestSuite::new(1, ApiVersion::V2);
    let key1 = txn_file_key(21);
    let key2 = txn_file_key(22);
    let key3 = txn_file_key(23);
    let index_key = txn_file_index_key(1);
    let region = get_ready_region(&mut suite, &key1);
    let region_id = region.id;

    let checkpoint_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let (mut req_tx, _event_feed_wrap, receive_event) =
        new_event_feed_v2(suite.get_region_cdc_client(region_id));
    let mut req = suite.new_changedata_request(region_id);
    req.request_id = 3;
    req.checkpoint_ts = checkpoint_ts.into_inner();
    req.set_start_key(Key::from_raw(b"").into_encoded());
    req.set_end_key(Key::from_raw(MAGIC_KEYSPACE_END).into_encoded());
    subscribe_until_initialized(
        &mut req_tx,
        receive_event.as_ref(),
        &req,
        "test_txn_file_cdc_multi_chunk_split_batches_delete_and_filter_index_rows initial subscribe",
    );

    let start_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let chunks = vec![
        vec![
            (key1.clone(), OP_PUT, vec![b'a'; 20 * 1024]),
            (
                index_key.clone(),
                OP_PUT,
                b"index-value-is-filtered".to_vec(),
            ),
        ],
        vec![(key2.clone(), OP_PUT, vec![b'b'; 20 * 1024])],
        vec![(key3.clone(), OP_DELETE, txn_file_delete_value())],
    ];
    let chunk_ids =
        build_txn_file_chunks_with_ops(&get_dfs(&suite), start_ts.into_inner() + 1, &chunks);
    must_txn_file_prewrite(
        &mut suite,
        region_id,
        key1.clone(),
        start_ts.into_inner(),
        chunk_ids,
    );

    let commit_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    must_txn_file_commit(
        &mut suite,
        region_id,
        start_ts.into_inner(),
        commit_ts.into_inner(),
    );

    let batches = collect_entry_batches(receive_event.as_ref(), EventLogType::Committed, 3);
    assert!(batches.len() > 1, "expected event splitting: {:?}", batches);
    let commit_rows = batches.into_iter().flatten().collect::<Vec<_>>();
    assert_eq!(commit_rows[0].get_key(), key1.as_slice());
    assert_eq!(commit_rows[0].get_value(), vec![b'a'; 20 * 1024].as_slice());
    assert_eq!(commit_rows[0].get_op_type(), EventRowOpType::Put);
    assert_eq!(commit_rows[1].get_key(), key2.as_slice());
    assert_eq!(commit_rows[1].get_value(), vec![b'b'; 20 * 1024].as_slice());
    assert_eq!(commit_rows[1].get_op_type(), EventRowOpType::Put);
    assert_eq!(commit_rows[2].get_key(), key3.as_slice());
    assert_eq!(commit_rows[2].get_value(), b"");
    assert_eq!(commit_rows[2].get_op_type(), EventRowOpType::Delete);
    assert!(
        commit_rows
            .iter()
            .all(|row| row.get_key() != index_key.as_slice())
    );
    for row in &commit_rows {
        assert_eq!(row.get_start_ts(), start_ts.into_inner(), "{:?}", row);
        assert_eq!(row.get_commit_ts(), commit_ts.into_inner(), "{:?}", row);
    }

    assert_resolved_after(receive_event.as_ref(), region_id, commit_ts.into_inner());
    block_on(req_tx.close()).unwrap();
    sleep_ms(20);
    suite.stop();
}

#[test]
fn test_txn_file_cdc_split_after_txn_file_commit_re_registers_without_row_loss() {
    let mut suite = TestSuite::new(1, ApiVersion::V2);
    let key1 = txn_file_key(31);
    let key2 = txn_file_key(32);
    let split_key = txn_file_key(32);
    let region = get_ready_region(&mut suite, &key1);
    let region_id = region.id;

    let checkpoint_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let (mut req_tx, _event_feed_wrap, receive_event) =
        new_event_feed_v2(suite.get_region_cdc_client(region_id));
    let mut req = suite.new_changedata_request(region_id);
    req.request_id = 4;
    req.checkpoint_ts = checkpoint_ts.into_inner();
    req.set_start_key(Key::from_raw(MAGIC_KEYSPACE_START).into_encoded());
    req.set_end_key(Key::from_raw(MAGIC_KEYSPACE_END).into_encoded());
    subscribe_until_initialized(
        &mut req_tx,
        receive_event.as_ref(),
        &req,
        "test_txn_file_cdc_split_after_txn_file_commit_re_registers_without_row_loss initial subscribe",
    );

    let start_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let rows = vec![
        (key1.clone(), txn_file_value(31)),
        (key2.clone(), txn_file_value(32)),
    ];
    let chunk_ids = build_txn_file_chunk(&get_dfs(&suite), start_ts.into_inner() + 1, &rows);
    must_txn_file_prewrite(
        &mut suite,
        region_id,
        key1.clone(),
        start_ts.into_inner(),
        chunk_ids,
    );

    let commit_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    must_txn_file_commit(
        &mut suite,
        region_id,
        start_ts.into_inner(),
        commit_ts.into_inner(),
    );

    suite.cluster.must_split(&region, &split_key);

    let (commit_rows, split_regions) =
        collect_rows_until_epoch_not_match(receive_event.as_ref(), EventLogType::Committed, 2);
    assert_eq!(commit_rows[0].get_key(), key1.as_slice());
    assert_eq!(commit_rows[1].get_key(), key2.as_slice());

    for region in &split_regions {
        suite.cluster.wait_region_info_ready(region);
        let mut split_req = ChangeDataRequest {
            region_id: region.get_id(),
            ..Default::default()
        };
        split_req
            .mut_header()
            .set_cluster_id(suite.cluster.pd_client().get_cluster_id().unwrap());
        split_req.request_id = 40 + region.get_id();
        split_req.checkpoint_ts = commit_ts.into_inner();
        split_req.set_region_epoch(region.get_region_epoch().clone());
        split_req.set_start_key(Key::from_raw(MAGIC_KEYSPACE_START).into_encoded());
        split_req.set_end_key(Key::from_raw(MAGIC_KEYSPACE_END).into_encoded());
        subscribe_until_initialized(
            &mut req_tx,
            receive_event.as_ref(),
            &split_req,
            "test_txn_file_cdc_split_after_txn_file_commit_re_registers_without_row_loss resubscribe",
        );
    }

    block_on(req_tx.close()).unwrap();
    sleep_ms(20);
    suite.stop();
}