// Copyright 2025 PingCAP, Inc.

use std::{collections::HashMap, fmt};

use api_version::ApiV2;
use bytes::Buf;
use cloud_encryption::ToEncryptionState;
use kvengine::{
    IdVer, IoType, Iterator as KvIterator, LOCK_CF, ShardTag, SnapAccess, TXN_FILE_REF, UserMeta,
    WRITE_CF, WriteBatch,
    table::{
        memtable::WriteBatchEntry,
        txn_file::{SkipOpTxnFileIterator, TxnFileIterator},
    },
    util::TxnFileRefExt,
};
use kvenginepb::TxnFileRefs;
use kvproto::{cdcpb, raft_cmdpb::AdminRequest};
use log_wrappers::Value as LogValue;
use protobuf::Message;
use rfstore::store::ApplyObserver;
use tidb_query_datatype::codec::table::{INDEX_PREFIX_SEP, PREFIX_LEN, TABLE_PREFIX};
use tikv_util::{debug, error, info, mpsc::Sender, trace};
use tokio::task::JoinHandle;
use txn_types::{LockType, TimeStamp, parse_lock};

use crate::{CdcMsg, metrics::REP_TXN_FILE_COUNTER_VEC};

pub struct CdcApplyObserver {
    store_id: u64,
    kv: kvengine::Engine,
    sender: Sender<CdcMsg>,
    runtime: tokio::runtime::Handle,
    join_handles: HashMap<u64 /* region_id */, Vec<JoinHandle<EventBuilder>>>,
}

#[derive(Default)]
pub struct RegionEvents {
    pub region_version: u64,
    pub events: Vec<cdcpb::Event>,
    pub tracked_locks: Vec<(Vec<u8>, u64)>,
}

impl fmt::Debug for RegionEvents {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RegionEvents")
            .field("region_ver", &self.region_version)
            .field("events", &self.events.len())
            .field("tracked_locks", &self.tracked_locks.len())
            .finish()
    }
}

impl CdcApplyObserver {
    pub fn new(
        kv: kvengine::Engine,
        sender: Sender<CdcMsg>,
        runtime: tokio::runtime::Handle,
    ) -> Self {
        let store_id = kv.get_engine_id();
        Self {
            store_id,
            kv,
            sender,
            runtime,
            join_handles: HashMap::default(),
        }
    }

    fn send_applied_msg(sender: &Sender<CdcMsg>, region_id: u64, region_events: RegionEvents) {
        let msg = CdcMsg::Applied {
            region_id,
            region_events,
        };
        if let Err(e) = sender.send(msg) {
            // Should happen only on stop.
            error!("send cdc event failed"; "region" => region_id, "error" => ?e);
        }
    }

    async fn flush_region_async(&mut self, region_id: u64) {
        if let Some(handles) = self.join_handles.remove(&region_id) {
            Self::flush_region_impl(region_id, handles, self.sender.clone()).await;
        };
    }

    async fn flush_region_impl(
        region_id: u64,
        handles: Vec<JoinHandle<EventBuilder>>,
        sender: Sender<CdcMsg>,
    ) {
        debug_assert!(!handles.is_empty());

        let mut events = RegionEvents::default();
        for (idx, handle) in handles.into_iter().enumerate() {
            // It's verbose to handle the task cancelled here.
            // Make sure the replication worker is stopped before tokio runtime shutdown.
            let mut event_builder = handle.await.expect("task panic/cancelled");

            if idx == 0 {
                events.region_version = event_builder.region_version();
            } else {
                debug_assert_eq!(
                    events.region_version,
                    event_builder.region_version(),
                    "region version changed, region_id: {}",
                    region_id,
                );
            }

            events.events.push(event_builder.build_event());
            events.tracked_locks.extend(event_builder.tracked_locks);
        }
        // Note: `sender` should be unbounded.
        Self::send_applied_msg(&sender, region_id, events);
    }

    async fn flush_all(&mut self) {
        let mut join_set = tokio::task::JoinSet::new();
        for (region_id, handles) in self.join_handles.drain() {
            let sender = self.sender.clone();
            join_set.spawn_on(
                async move {
                    Self::flush_region_impl(region_id, handles, sender).await;
                },
                &self.runtime,
            );
        }
        while let Some(res) = join_set.join_next().await {
            // It's verbose to handle the task cancelled here.
            // Make sure the replication worker is stopped before tokio runtime shutdown.
            res.expect("task panic/cancelled");
        }
    }
}

impl ApplyObserver for CdcApplyObserver {
    fn on_apply(&mut self, region_id: u64, region_version: u64, log_index: u64, wb: &WriteBatch) {
        if let Some(txn_file_refs_bin) = wb.get_property(TXN_FILE_REF) {
            debug_assert_eq!(wb.cf_len(WRITE_CF) + wb.cf_len(LOCK_CF), 0);
            self.on_apply_txn_file(
                region_id,
                region_version,
                log_index,
                txn_file_refs_bin.chunk(),
            );
        }

        if wb.cf_len(WRITE_CF) + wb.cf_len(LOCK_CF) == 0 {
            return;
        }

        let write_cf = wb.get_cf(WRITE_CF);
        let snap_access = self.kv.get_snap_access(region_id).unwrap();
        debug_assert_eq!(snap_access.get_version(), region_version);
        let mut event_builder = EventBuilder::new(snap_access, log_index);
        write_cf.iterate(|entry, buf, value| {
            event_builder.add_row(entry, buf, value);
        });
        let lock_cf = wb.get_cf(LOCK_CF);
        lock_cf.iterate(|entry, buf, value| {
            event_builder.add_lock(entry, buf, value);
        });
        let handle = self.runtime.spawn(async move {
            event_builder.fetch_old_values().await;
            event_builder
        });
        self.join_handles.entry(region_id).or_default().push(handle);
    }

    fn on_apply_admin(
        &mut self,
        region_id: u64,
        region_version: u64,
        log_index: u64,
        admin: &AdminRequest,
    ) {
        if admin.has_splits()
            || admin.has_prepare_merge()
            || admin.has_commit_merge()
            || admin.has_rollback_merge()
        {
            let tag = ShardTag::new(self.store_id, IdVer::new(region_id, region_version));
            self.flush_region(region_id);
            info!("{} on apply admin {:?}", tag, admin; "log_index" => log_index);
            let msg = CdcMsg::AppliedAdmin {
                region_id,
                region_version,
                admin: admin.clone(),
            };
            if let Err(e) = self.sender.send(msg) {
                error!("{} send cdc event failed", tag; "error" => ?e);
            }
        }
    }

    fn flush(&mut self) {
        self.runtime.clone().block_on(self.flush_all());
    }

    fn flush_region(&mut self, region_id: u64) {
        self.runtime
            .clone()
            .block_on(self.flush_region_async(region_id));
    }
}

struct EventBuilder {
    snap_access: SnapAccess,
    index: u64,
    rows: cdcpb::EventEntries,
    tracked_locks: Vec<(Vec<u8>, u64)>,
}

impl EventBuilder {
    fn new(snap_access: SnapAccess, index: u64) -> Self {
        Self {
            snap_access,
            index,
            rows: cdcpb::EventEntries::default(),
            tracked_locks: Vec::new(),
        }
    }

    fn region_version(&self) -> u64 {
        self.snap_access.get_version()
    }

    #[inline]
    fn tag(&self) -> ShardTag {
        self.snap_access.get_tag()
    }

    fn add_row(&mut self, entry: &WriteBatchEntry, buf: &[u8], value: &[u8]) {
        let entry_key = entry.key(buf);
        if is_index_key(entry_key) {
            // Skip index keys as CDC does not need them.
            return;
        }
        let mut event_row = cdcpb::EventRow::default();
        event_row.set_key(entry_key.to_vec());
        event_row.set_value(value.to_vec());
        event_row.set_commit_ts(entry.version);
        let user_meta = UserMeta::from_slice(entry.user_meta(buf));
        event_row.set_start_ts(user_meta.start_ts);
        if event_row.get_value().is_empty() {
            event_row.set_op_type(cdcpb::EventRowOpType::Delete);
        } else {
            event_row.set_op_type(cdcpb::EventRowOpType::Put);
        }
        event_row.set_type(cdcpb::EventLogType::Committed);
        self.rows.mut_entries().push(event_row);

        trace!("{} add_row", self.tag();
            "key" => LogValue::key(entry_key),
            "version" => entry.version,
            "start_ts" => user_meta.start_ts,
            "commit_ts" => user_meta.commit_ts,
            "log_index" => self.index,
        );
    }

    fn add_lock(&mut self, entry: &WriteBatchEntry, buf: &[u8], value: &[u8]) {
        let entry_key = entry.key(buf);
        if is_index_key(entry_key) {
            // Skip index keys as CDC does not need them.
            return;
        }
        if value.is_empty() {
            self.tracked_locks.push((entry_key.to_vec(), 0));

            trace!("{} add_lock (untrack)", self.tag();
                "key" => LogValue::key(entry_key),
                "version" => entry.version,
                "log_index" => self.index,
            );
            return;
        }
        let lock = match parse_lock(value).unwrap() {
            tikv_util::Either::Left(lock) => lock,
            tikv_util::Either::Right(shared_locks) => {
                trace!("{} add_lock (skip shared)", self.tag();
                    "key" => LogValue::key(entry_key),
                    "version" => entry.version,
                    "shared_holders" => shared_locks.len(),
                    "log_index" => self.index,
                );
                return;
            }
        };
        if !matches!(lock.lock_type, LockType::Put | LockType::Delete) {
            // Note: other lock types is not tracked, but will still be untracked.
            // The resolver can handle this case.
            trace!("{} add_lock (skip)", self.tag();
                "key" => LogValue::key(entry_key),
                "version" => entry.version,
                "lock.ts" => lock.ts,
                "lock.ty" => ?lock.lock_type,
                "log_index" => self.index,
            );
            return;
        }
        self.tracked_locks
            .push((entry_key.to_vec(), lock.ts.into_inner()));

        trace!("{} add_lock", self.tag();
            "key" => LogValue::key(entry_key),
            "version" => entry.version,
            "lock.ts" => lock.ts,
            "lock.ty" => ?lock.lock_type,
            "log_index" => self.index,
        );
    }

    fn add_txn_file_row(&mut self, key: &[u8], value: kvengine::table::Value, start_ts: u64) {
        let value_bytes = value.get_value();

        let mut event_row = cdcpb::EventRow::default();
        event_row.set_key(key.to_vec());
        event_row.set_value(value_bytes.to_vec());
        event_row.set_commit_ts(value.version);
        event_row.set_start_ts(start_ts);
        if value_bytes.is_empty() {
            event_row.set_op_type(cdcpb::EventRowOpType::Delete);
        } else {
            event_row.set_op_type(cdcpb::EventRowOpType::Put);
        }
        event_row.set_type(cdcpb::EventLogType::Committed);
        self.rows.mut_entries().push(event_row);

        trace!("{} add_txn_file_row", self.tag();
            "key" => LogValue::key(key),
            "commit_ts" => value.version,
            "start_ts" => start_ts,
            "log_index" => self.index,
        );
    }

    #[inline]
    fn track_txn_file_lock(&mut self, key: &[u8], lock_ts: TimeStamp) {
        self.tracked_locks
            .push((key.to_vec(), lock_ts.into_inner()));

        trace!("{} track_txn_file_lock", self.tag();
            "key" => LogValue::key(key),
            "lock_ts" => lock_ts,
            "log_index" => self.index,
        );
    }

    #[inline]
    fn untrack_txn_file_lock(&mut self, key: &[u8], start_ts: u64) {
        self.tracked_locks.push((key.to_vec(), 0));

        trace!("{} untrack_txn_file_lock", self.tag();
            "key" => LogValue::key(key),
            "start_ts" => start_ts,
            "log_index" => self.index,
        );
    }

    async fn fetch_old_values(&mut self) {
        for row in self.rows.entries.iter_mut() {
            if row.get_type() != cdcpb::EventLogType::Committed {
                continue;
            }
            let mut keyspace_key =
                ApiV2::get_keyspace_prefix_by_id(self.snap_access.get_keyspace_id());
            keyspace_key.extend_from_slice(row.get_key());
            let old_item = self
                .snap_access
                .get_async(WRITE_CF, &keyspace_key, row.commit_ts - 1)
                .await;
            if !old_item.get_value().is_empty() {
                row.set_old_value(old_item.get_value().to_vec());
            }
        }
    }

    fn build_event(&mut self) -> cdcpb::Event {
        let mut event = cdcpb::Event::default();
        event.index = self.index;
        event.region_id = self.snap_access.get_id();
        event.mut_entries().set_entries(self.rows.take_entries());
        event
    }
}

impl CdcApplyObserver {
    // Panic on error. Otherwise, events may be lost.
    // TODO: Blacklist the keyspace.
    fn on_apply_txn_file(
        &mut self,
        region_id: u64,
        region_version: u64,
        log_index: u64,
        txn_file_refs_bin: &[u8],
    ) {
        let snap_access = self.kv.get_snap_access(region_id).unwrap();
        assert_eq!(snap_access.get_version(), region_version);

        let tag = snap_access.get_tag();

        let mut txn_file_refs = TxnFileRefs::new();
        txn_file_refs
            .merge_from_bytes(txn_file_refs_bin)
            .unwrap_or_else(|err| panic!("{} failed to unmarshal TxnFileRefs: {:?}", tag, err));
        assert_eq!(
            txn_file_refs.txn_file_refs.len(),
            1,
            "{} expected exactly one TxnFileRef",
            tag
        );
        let txn_file_ref = txn_file_refs.take_txn_file_refs().pop().unwrap();
        debug!("{} on_apply_txn_file {:?}", tag, txn_file_ref; "log_index" => log_index);

        let is_locked = txn_file_ref.is_locked();
        let lock_prefix = is_locked.then(|| {
            txn_types::parse_non_shared_lock(txn_file_ref.get_lock_val_prefix())
                .unwrap_or_else(|err| panic!("{} failed to parse lock_val_prefix: {:?}", tag, err))
        });

        let (is_commit, is_rollback) = txn_file_ref.get_commit_or_rollback();
        let start_ts = (is_commit || is_rollback)
            .then(|| UserMeta::from_slice(txn_file_ref.get_user_meta()).start_ts);

        let txn_chunk_manager = self.kv.get_txn_chunk_manager();
        let encryption_state = snap_access.get_encryption_key().to_encryption_state();
        let runtime = self.runtime.clone();

        let mut event_builder = EventBuilder::new(snap_access, log_index);
        let task = async move {
            let prepare_task = runtime.spawn_blocking(move || {
                // Set `is_prepared` false to handle the case when apply is slow and the
                // prepared txn file is already GCed.
                let is_prepared = false;
                txn_chunk_manager.load_txn_file_from_ref(
                    region_id,
                    region_version,
                    &txn_file_ref,
                    is_prepared,
                    encryption_state,
                    IoType::Export,
                )
            });
            let txn_file = prepare_task
                .await
                .expect("task panic/cancelled")
                .unwrap_or_else(|err| panic!("{} load_txn_file_from_ref failed: {:?}", tag, err));

            let reverse = false;
            let txn_file_iter = TxnFileIterator::new(txn_file, reverse);
            let mut iter = SkipOpTxnFileIterator::new(
                txn_file_iter,
                true, // skip_lock, we need PUT/DELETE only.
                true, // skip_check_not_exist
            );
            iter.rewind();

            if is_locked {
                let lock_ts = lock_prefix.unwrap().ts;
                while iter.valid() {
                    let key = iter.key();
                    if !is_index_key(&key) {
                        event_builder.track_txn_file_lock(&key, lock_ts);
                    }
                    iter.next();
                }
                REP_TXN_FILE_COUNTER_VEC.lock.inc();
            } else if is_commit {
                let start_ts = start_ts.unwrap();
                while iter.valid() {
                    let key = iter.key();
                    if !is_index_key(&key) {
                        let value = iter.value();
                        event_builder.add_txn_file_row(&key, value, start_ts);
                        event_builder.untrack_txn_file_lock(&key, start_ts);
                    }
                    iter.next();
                }
                event_builder.fetch_old_values().await;
                REP_TXN_FILE_COUNTER_VEC.commit.inc();
            } else if is_rollback {
                let start_ts = start_ts.unwrap();
                while iter.valid() {
                    let key = iter.key();
                    if !is_index_key(&key) {
                        event_builder.untrack_txn_file_lock(&key, start_ts);
                    }
                    iter.next();
                }
                REP_TXN_FILE_COUNTER_VEC.rollback.inc();
            } else {
                unreachable!()
            }

            event_builder
        };

        let handle = self.runtime.spawn(task);
        self.join_handles.entry(region_id).or_default().push(handle);
    }
}

pub(crate) fn is_index_key(key: &[u8]) -> bool {
    if key.len() < PREFIX_LEN {
        return false;
    }
    let trimmed_key = &key[..PREFIX_LEN];
    trimmed_key.starts_with(TABLE_PREFIX) && trimmed_key.ends_with(INDEX_PREFIX_SEP)
}

#[cfg(test)]
mod tests {
    use api_version::ApiV2;
    use kvengine::test_engine;
    use txn_types::{Lock, SharedLocks};

    use super::*;

    #[test]
    fn test_add_lock_tracks() {
        let (engine, _apply_tx) = test_engine::new_test_engine_api_v2();
        let snap_access = engine.get_snap_access(1).unwrap();
        let mut event_builder = EventBuilder::new(snap_access, 1);

        let keyspace_prefix = ApiV2::get_keyspace_prefix_by_id(1);
        let mut outer_key_put = keyspace_prefix.clone();
        outer_key_put.extend_from_slice(b"k_put");
        let mut outer_key_lock = keyspace_prefix.clone();
        outer_key_lock.extend_from_slice(b"k_lock");
        let mut outer_key_pess = keyspace_prefix.clone();
        outer_key_pess.extend_from_slice(b"k_pess");
        let mut outer_key_del = keyspace_prefix;
        outer_key_del.extend_from_slice(b"k_del");

        let lock_put_ts = 10;
        let lock_del_ts = 20;
        let lock_lock_ts = 30;
        let lock_pess_ts = 40;

        let put_lock = Lock::new(
            LockType::Put,
            b"pk".to_vec(),
            lock_put_ts.into(),
            0,
            None,
            TimeStamp::zero(),
            0,
            TimeStamp::zero(),
        )
        .to_bytes();
        let lock_only_lock = Lock::new(
            LockType::Lock,
            b"pk".to_vec(),
            lock_lock_ts.into(),
            0,
            None,
            TimeStamp::zero(),
            0,
            TimeStamp::zero(),
        )
        .to_bytes();
        let pessimistic_lock = Lock::new(
            LockType::Pessimistic,
            b"pk".to_vec(),
            lock_pess_ts.into(),
            0,
            None,
            TimeStamp::zero(),
            0,
            TimeStamp::zero(),
        )
        .to_bytes();
        let delete_lock = Lock::new(
            LockType::Delete,
            b"pk".to_vec(),
            lock_del_ts.into(),
            0,
            None,
            TimeStamp::zero(),
            0,
            TimeStamp::zero(),
        )
        .to_bytes();

        let mut wb = WriteBatch::default();
        wb.reset(1);
        wb.put(LOCK_CF, &outer_key_put, &put_lock, 0, &[], 0);
        wb.put(LOCK_CF, &outer_key_lock, &lock_only_lock, 0, &[], 0);
        wb.put(LOCK_CF, &outer_key_pess, &pessimistic_lock, 0, &[], 0);
        wb.put(LOCK_CF, &outer_key_del, &delete_lock, 0, &[], 0);

        let lock_cf = wb.get_cf(LOCK_CF);
        lock_cf.iterate(|entry, buf, value| {
            event_builder.add_lock(entry, buf, value);
        });

        assert_eq!(
            event_builder.tracked_locks,
            vec![
                (b"k_put".to_vec(), lock_put_ts),
                (b"k_del".to_vec(), lock_del_ts),
            ]
        );
    }

    #[test]
    fn test_add_lock_skips_shared_lock_wrapper() {
        let (engine, _apply_tx) = test_engine::new_test_engine_api_v2();
        let snap_access = engine.get_snap_access(1).unwrap();
        let mut event_builder = EventBuilder::new(snap_access, 1);

        let mut outer_key = ApiV2::get_keyspace_prefix_by_id(1);
        outer_key.extend_from_slice(b"k_shared");

        let mut shared_locks = SharedLocks::new();
        shared_locks
            .insert_lock(Lock::new(
                LockType::Lock,
                b"pk1".to_vec(),
                10.into(),
                0,
                None,
                TimeStamp::zero(),
                0,
                TimeStamp::zero(),
            ))
            .unwrap();
        shared_locks
            .insert_lock(Lock::new(
                LockType::Lock,
                b"pk2".to_vec(),
                20.into(),
                0,
                None,
                TimeStamp::zero(),
                0,
                TimeStamp::zero(),
            ))
            .unwrap();

        let mut wb = WriteBatch::default();
        wb.reset(1);
        wb.put(LOCK_CF, &outer_key, &shared_locks.to_bytes(), 0, &[], 0);

        let lock_cf = wb.get_cf(LOCK_CF);
        lock_cf.iterate(|entry, buf, value| {
            event_builder.add_lock(entry, buf, value);
        });

        assert!(event_builder.tracked_locks.is_empty());
    }

    #[test]
    fn test_add_lock_untracks() {
        let (engine, _apply_tx) = test_engine::new_test_engine_api_v2();
        let snap_access = engine.get_snap_access(1).unwrap();
        let mut event_builder = EventBuilder::new(snap_access, 1);

        let mut outer_key = ApiV2::get_keyspace_prefix_by_id(1);
        outer_key.extend_from_slice(b"k");

        let lock_only_lock = Lock::new(
            LockType::Lock,
            b"pk".to_vec(),
            10.into(),
            0,
            None,
            TimeStamp::zero(),
            0,
            TimeStamp::zero(),
        )
        .to_bytes();

        let mut wb = WriteBatch::default();
        wb.reset(1);
        wb.put(LOCK_CF, &outer_key, &lock_only_lock, 0, &[], 0);
        wb.delete(LOCK_CF, &outer_key, 0);

        let lock_cf = wb.get_cf(LOCK_CF);
        lock_cf.iterate(|entry, buf, value| {
            event_builder.add_lock(entry, buf, value);
        });

        assert_eq!(event_builder.tracked_locks, vec![(b"k".to_vec(), 0)]);
    }
}