// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2025 PingCAP, Inc.

use std::{
    sync::Arc,
    thread,
    time::{Duration, Instant},
};

use api_version::{KvFormat, test_kv_format_impl};
use futures::{executor::block_on, sink::SinkExt};
use grpcio::{ClientDuplexSender, WriteFlags};
use kvproto::{cdcpb::*, kvrpcpb::*, metapb::RegionEpoch};
use pd_client::PdClient;
use raft::StateRole;
use raftstore::coprocessor::{ObserverContext, RoleChange, RoleObserver};
use test_raftstore::sleep_ms;

use crate::{TestSuite, get_ready_region, new_event_feed, subscribe_until_initialized};

const EPOCH_NOT_MATCH_TIMEOUT: Duration = Duration::from_secs(15);

#[derive(Debug, PartialEq)]
enum EpochNotMatchBatchResult {
    Continue,
    Retry,
    Done,
}

fn validate_epoch_not_match_batch(
    events: Vec<Event>,
    initialized: &mut bool,
) -> EpochNotMatchBatchResult {
    assert!(!events.is_empty());
    let mut result = EpochNotMatchBatchResult::Continue;
    let mut terminal_error = false;
    for event in events {
        assert!(!terminal_error, "event after terminal error: {:?}", event);
        match event.event.unwrap() {
            Event_oneof_event::Error(err) if err.has_epoch_not_match() => {
                terminal_error = true;
                result = EpochNotMatchBatchResult::Done;
            }
            Event_oneof_event::Error(err) if err.has_region_not_found() || err.has_not_leader() => {
                terminal_error = true;
                result = EpochNotMatchBatchResult::Retry;
            }
            Event_oneof_event::Entries(es) => {
                assert_eq!(es.entries.len(), 1, "{:?}", es);
                assert_eq!(
                    es.entries[0].get_type(),
                    EventLogType::Initialized,
                    "{:?}",
                    es
                );
                assert!(!*initialized, "duplicate initialized event: {:?}", es);
                *initialized = true;
            }
            other => panic!("unknown event {:?}", other),
        }
    }
    result
}

fn wait_for_epoch_not_match(
    receive_event: &mut dyn FnMut() -> ChangeDataEvent,
    retry: &mut dyn FnMut(),
    timeout: Duration,
) {
    let timer = Instant::now();
    let mut initialized = false;
    loop {
        assert!(
            timer.elapsed() <= timeout,
            "timed out waiting for epoch not match"
        );
        let events = receive_event().events.to_vec();
        assert!(
            timer.elapsed() <= timeout,
            "timed out waiting for epoch not match"
        );
        match validate_epoch_not_match_batch(events, &mut initialized) {
            EpochNotMatchBatchResult::Continue => {}
            EpochNotMatchBatchResult::Retry => retry(),
            EpochNotMatchBatchResult::Done => return,
        }
    }
}

// Assert epoch not match error is received, with retry on region not found or
// not leader error, and accept optional initialized event.
fn assert_epoch_not_match_with_retry(
    req_tx: &mut ClientDuplexSender<ChangeDataRequest>,
    receive_event: &dyn Fn(bool) -> ChangeDataEvent,
    req: &ChangeDataRequest,
) {
    wait_for_epoch_not_match(
        &mut || receive_event(false),
        &mut || {
            sleep_ms(20);
            block_on(req_tx.send((req.clone(), WriteFlags::default()))).unwrap();
        },
        EPOCH_NOT_MATCH_TIMEOUT,
    );
}

#[test]
fn test_failed_pending_batch() {
    // For test that a pending cmd batch contains a error like epoch not match.
    let mut suite = TestSuite::new(1, ApiVersion::V2);

    let fp = "cdc_incremental_scan_start";
    fail::cfg(fp, "pause").unwrap();

    let region = get_ready_region(&mut suite, &[]);
    let mut req = suite.new_changedata_request(region.get_id());
    let (mut req_tx, event_feed_wrap, receive_event) =
        new_event_feed(suite.get_region_cdc_client(region.id));
    block_on(req_tx.send((req.clone(), WriteFlags::default()))).unwrap();
    // Split region.
    suite.cluster.must_split(&region, b"xk0");
    // Wait for receiving split cmd.
    sleep_ms(200);
    fail::remove(fp);

    loop {
        let mut events = receive_event(false).events.to_vec();
        match events.pop().unwrap().event.unwrap() {
            Event_oneof_event::Error(err) => {
                assert!(err.has_epoch_not_match(), "{:?}", err);
                break;
            }
            Event_oneof_event::Entries(es) => {
                assert!(es.entries.len() == 1, "{:?}", es);
                let e = &es.entries[0];
                assert_eq!(e.get_type(), EventLogType::Initialized, "{:?}", es);
            }
            other => panic!("unknown event {:?}", other),
        }
    }
    // Try to subscribe region again.
    // next-gen: split generates new regions at left side.
    let region = get_ready_region(&mut suite, b"xk1");
    // Ensure it is the previous region.
    assert_eq!(req.get_region_id(), region.get_id());
    req.set_region_epoch(region.get_region_epoch().clone());
    subscribe_until_initialized(&mut req_tx, receive_event.as_ref(), &req, "after split");

    event_feed_wrap.replace(None);
    suite.stop();
}

#[test]
fn test_region_ready_after_deregister() {
    test_kv_format_impl!(test_region_ready_after_deregister_impl<ApiV2>);
}

fn test_region_ready_after_deregister_impl<F: KvFormat>() {
    let mut suite = TestSuite::new(1, F::TAG);

    let fp = "cdc_incremental_scan_start";
    fail::cfg(fp, "pause").unwrap();

    let region = get_ready_region(&mut suite, b"");
    let req = suite.new_changedata_request(region.id);
    let (mut req_tx, event_feed_wrap, receive_event) =
        new_event_feed(suite.get_region_cdc_client(region.id));
    block_on(req_tx.send((req, WriteFlags::default()))).unwrap();
    // Sleep for a while to make sure the region has been subscribed
    sleep_ms(200);

    // Simulate a role change event
    let region = get_ready_region(&mut suite, &[]);
    let leader = suite.cluster.leader_of_region(region.get_id()).unwrap();
    let mut context = ObserverContext::new(&region);
    suite
        .obs
        .get(&leader.get_store_id())
        .unwrap()
        .on_role_change(&mut context, &RoleChange::new(StateRole::Follower));

    // Then CDC should not panic
    fail::remove(fp);
    receive_event(false);

    event_feed_wrap.replace(None);
    suite.stop();
}

#[test]
fn test_connections_register() {
    test_kv_format_impl!(test_connections_register_impl<ApiV2>);
}

fn test_connections_register_impl<F: KvFormat>() {
    let mut suite = TestSuite::new(1, F::TAG);

    let fp = "cdc_incremental_scan_start";
    fail::cfg(fp, "pause").unwrap();

    let (k, v) = ("xkey1".to_owned(), "value".to_owned());
    // Region info
    let region = get_ready_region(&mut suite, &[]);
    // Prewrite
    let start_ts = block_on(suite.cluster.pd_client().get_tso()).unwrap();
    let mut mutation = Mutation::default();
    mutation.set_op(Op::Put);
    mutation.key = k.clone().into_bytes();
    mutation.value = v.into_bytes();
    suite.must_kv_prewrite(region.get_id(), vec![mutation], k.into_bytes(), start_ts);

    let mut req = suite.new_changedata_request(region.get_id());
    req.set_region_epoch(RegionEpoch::default());

    let (mut req_tx, event_feed_wrap, receive_event) =
        new_event_feed(suite.get_region_cdc_client(region.id));
    block_on(req_tx.send((req.clone(), WriteFlags::default()))).unwrap();
    let mut events = receive_event(false).events.to_vec();
    match events.pop().unwrap().event.unwrap() {
        Event_oneof_event::Error(err) => {
            assert!(err.has_epoch_not_match(), "{:?}", err);
        }
        other => panic!("unknown event {:?}", other),
    }

    // Conn 1
    req.set_region_epoch(region.get_region_epoch().clone());
    block_on(req_tx.send((req.clone(), WriteFlags::default()))).unwrap();
    thread::sleep(Duration::from_secs(1));
    // Close conn 1
    event_feed_wrap.replace(None);
    // Conn 2
    let (mut req_tx, resp_rx) = suite
        .get_region_cdc_client(region.get_id())
        .event_feed()
        .unwrap();
    block_on(req_tx.send((req.clone(), WriteFlags::default()))).unwrap();
    event_feed_wrap.replace(Some(resp_rx));
    // Split region.
    suite.cluster.must_split(&region, b"xk0");
    fail::remove(fp);
    // The split happens before the pause failpoint is removed, so the pending
    // subscription must eventually receive the authoritative epoch mismatch.
    assert_epoch_not_match_with_retry(&mut req_tx, receive_event.as_ref(), &req);

    event_feed_wrap.replace(None);
    suite.stop();
}

#[test]
fn test_merge() {
    test_kv_format_impl!(test_merge_impl<ApiV2>);
}

fn test_merge_impl<F: KvFormat>() {
    let mut suite = TestSuite::new(1, F::TAG);
    // Split region
    let region = suite.cluster.get_region(&[]);
    suite.cluster.must_split(&region, b"xk1");
    // Subscribe source region
    let source = get_ready_region(&mut suite, b"xk0");
    let mut req = suite.new_changedata_request(region.get_id());
    req.region_id = source.get_id();
    req.set_region_epoch(source.get_region_epoch().clone());
    let (mut source_tx, source_wrap, source_event) =
        new_event_feed(suite.get_region_cdc_client(source.get_id()));
    subscribe_until_initialized(
        &mut source_tx,
        source_event.as_ref(),
        &req,
        "test_merge source subscribe",
    );
    // Subscribe target region
    let target = get_ready_region(&mut suite, b"xk2");
    req.region_id = target.get_id();
    req.set_region_epoch(target.get_region_epoch().clone());
    let (mut target_tx, target_wrap, target_event) =
        new_event_feed(suite.get_region_cdc_client(target.get_id()));
    subscribe_until_initialized(
        &mut target_tx,
        target_event.as_ref(),
        &req,
        "test_merge target subscribe",
    );
    sleep_ms(200);
    // Pause before completing commit merge
    let commit_merge_fp = "before_handle_catch_up_logs_for_merge";
    fail::cfg(commit_merge_fp, "pause").unwrap();
    // The call is finished when prepare_merge is applied.
    suite.cluster.try_merge(source.get_id(), target.get_id());
    // Epoch not match after prepare_merge
    let mut events = source_event(false).events.to_vec();
    assert_eq!(events.len(), 1, "{:?}", events);
    match events.pop().unwrap().event.unwrap() {
        Event_oneof_event::Error(err) => {
            assert!(err.has_epoch_not_match(), "{:?}", err);
        }
        other => panic!("unknown event {:?}", other),
    }
    // Continue to commit merge
    let destroy_peer_fp = "destroy_peer";
    fail::cfg(destroy_peer_fp, "pause").unwrap();
    fail::remove(commit_merge_fp);
    // Wait until raftstore receives MergeResult
    sleep_ms(100);
    // Retry to subscribe source region
    let mut source_epoch = source.get_region_epoch().clone();
    source_epoch.set_version(source_epoch.get_version() + 1);
    source_epoch.set_conf_ver(source_epoch.get_conf_ver() + 1);
    req.region_id = source.get_id();
    req.set_region_epoch(source_epoch);
    block_on(source_tx.send((req, WriteFlags::default()))).unwrap();
    // Wait until raftstore receives ChangeCmd
    sleep_ms(100);
    fail::remove(destroy_peer_fp);
    loop {
        let mut events = source_event(false).events.to_vec();
        assert_eq!(events.len(), 1, "{:?}", events);
        match events.pop().unwrap().event.unwrap() {
            Event_oneof_event::Error(err) => {
                assert!(err.has_region_not_found(), "{:?}", err);
                break;
            }
            Event_oneof_event::Entries(es) => {
                assert!(es.entries.len() == 1, "{:?}", es);
                let e = &es.entries[0];
                assert_eq!(e.get_type(), EventLogType::Initialized, "{:?}", es);
            }
            other => panic!("unknown event {:?}", other),
        }
    }
    let mut events = target_event(false).events.to_vec();
    assert_eq!(events.len(), 1, "{:?}", events);
    match events.pop().unwrap().event.unwrap() {
        Event_oneof_event::Error(err) => {
            assert!(err.has_epoch_not_match(), "{:?}", err);
        }
        other => panic!("unknown event {:?}", other),
    }

    source_wrap.replace(None);
    target_wrap.replace(None);
    suite.stop();
}

#[test]
fn test_connections_registertest_deregister_pending_downstream() {
    test_kv_format_impl!(test_deregister_pending_downstream_impl<ApiV2>);
}

fn test_deregister_pending_downstream_impl<F: KvFormat>() {
    let mut suite = TestSuite::new(1, F::TAG);

    let build_resolver_fp = "before_schedule_resolver_ready";
    fail::cfg(build_resolver_fp, "pause").unwrap();
    let region = get_ready_region(&mut suite, b"");
    let node_id = suite.cluster.pd_client().get_all_stores(true).unwrap()[0].id;
    let mut req = suite.new_changedata_request(region.id);
    let (mut req_tx1, event_feed_wrap, receive_event) =
        new_event_feed(suite.get_region_cdc_client(region.id));
    block_on(req_tx1.send((req.clone(), WriteFlags::default()))).unwrap();
    suite.must_wait_delegate_condition(
        node_id,
        region.id,
        Arc::new(|d| {
            d.is_some_and(|d| {
                d.downstreams().len() == 1
                    && d.downstreams()[0]
                        .get_state()
                        .load()
                        .ready_for_change_events()
            })
        }),
    );

    let raft_capture_fp = "raft_on_capture_change";
    fail::cfg(raft_capture_fp, "pause").unwrap();

    // Conn 2
    let (mut req_tx2, resp_rx2) = suite.get_region_cdc_client(region.id).event_feed().unwrap();
    req.set_region_epoch(RegionEpoch::default());
    block_on(req_tx2.send((req.clone(), WriteFlags::default()))).unwrap();
    let _resp_rx1 = event_feed_wrap.replace(Some(resp_rx2));
    suite.must_wait_delegate_condition(
        node_id,
        region.id,
        Arc::new(|d| d.is_some_and(|d| d.downstreams().len() == 2)),
    );
    fail::remove(build_resolver_fp);
    let mut events = receive_event(false).events.to_vec();
    assert_eq!(events.len(), 1, "{:?}", events);
    match events.pop().unwrap().event.unwrap() {
        Event_oneof_event::Error(err) => {
            assert!(err.has_epoch_not_match(), "{:?}", err);
        }
        other => panic!("unknown event {:?}", other),
    }

    block_on(req_tx2.send((req, WriteFlags::default()))).unwrap();
    let mut events = receive_event(false).events.to_vec();
    assert_eq!(events.len(), 1, "{:?}", events);
    match events.pop().unwrap().event.unwrap() {
        Event_oneof_event::Error(err) => {
            assert!(err.has_epoch_not_match(), "{:?}", err);
        }
        other => panic!("unknown event {:?}", other),
    }
    fail::remove(raft_capture_fp);

    event_feed_wrap.replace(None);
    suite.stop();
}

#[cfg(test)]
mod tests {
    use std::{
        collections::VecDeque,
        panic::{AssertUnwindSafe, catch_unwind},
    };

    use kvproto::errorpb;

    use super::*;

    fn initialized_event() -> Event {
        let mut row = EventRow::default();
        row.set_type(EventLogType::Initialized);
        Event {
            event: Some(Event_oneof_event::Entries(EventEntries {
                entries: vec![row].into(),
                ..Default::default()
            })),
            ..Default::default()
        }
    }

    fn epoch_not_match_event() -> Event {
        let mut err = Error::default();
        err.set_epoch_not_match(errorpb::EpochNotMatch::default());
        Event {
            event: Some(Event_oneof_event::Error(err)),
            ..Default::default()
        }
    }

    fn not_leader_event() -> Event {
        let mut err = Error::default();
        err.set_not_leader(errorpb::NotLeader::default());
        Event {
            event: Some(Event_oneof_event::Error(err)),
            ..Default::default()
        }
    }

    fn change_data_event(events: Vec<Event>) -> ChangeDataEvent {
        ChangeDataEvent {
            events: events.into(),
            ..Default::default()
        }
    }

    #[test]
    fn test_epoch_not_match_validation() {
        let cases = [
            (
                "initialized before epoch not match",
                vec![initialized_event(), epoch_not_match_event()],
                Some(EpochNotMatchBatchResult::Done),
            ),
            (
                "initialized after epoch not match",
                vec![epoch_not_match_event(), initialized_event()],
                None,
            ),
            (
                "retryable error after epoch not match",
                vec![epoch_not_match_event(), not_leader_event()],
                None,
            ),
            (
                "initialized after retryable error",
                vec![not_leader_event(), initialized_event()],
                None,
            ),
            (
                "epoch not match after retryable error",
                vec![not_leader_event(), epoch_not_match_event()],
                None,
            ),
            (
                "duplicate initialized",
                vec![initialized_event(), initialized_event()],
                None,
            ),
            (
                "duplicate epoch not match",
                vec![epoch_not_match_event(), epoch_not_match_event()],
                None,
            ),
        ];
        for (name, events, expected) in cases {
            let result = catch_unwind(AssertUnwindSafe(|| {
                validate_epoch_not_match_batch(events, &mut false)
            }));
            match (result, expected) {
                (Ok(actual), Some(expected)) => assert_eq!(actual, expected, "case {name}"),
                (Err(_), None) => {}
                (Ok(actual), None) => {
                    panic!("case {name}: expected panic, got {actual:?}")
                }
                (Err(_), Some(_)) => panic!("case {name}: unexpected panic"),
            }
        }

        let mut responses = VecDeque::from([
            change_data_event(vec![initialized_event()]),
            change_data_event(vec![initialized_event()]),
        ]);
        let result = catch_unwind(AssertUnwindSafe(|| {
            wait_for_epoch_not_match(
                &mut || {
                    let response = responses.pop_front().unwrap();
                    if responses.is_empty() {
                        thread::sleep(Duration::from_millis(10));
                    }
                    response
                },
                &mut || {},
                Duration::from_millis(5),
            )
        }));
        if result.is_ok() {
            panic!("case global timeout: expected panic");
        }
    }
}