// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{fmt::Formatter, mem};

use byteorder::{ByteOrder, LittleEndian};
use bytes::{Buf, BufMut};
use kvengine::{IdVer, table::SnapVersion};
use kvenginepb::get_any_snap_from_changeset;
use kvproto::raft_cmdpb::{CustomRequest, RaftCmdRequest};
use protobuf::Message;
use tikv_util::{codec::number::U64_SIZE, warn};

use crate::store::{PeerTag, is_change_set_affect_mem_table};

pub fn get_custom_log(req: &RaftCmdRequest) -> Option<CustomRaftLog<'_>> {
    if !req.has_custom_request() {
        return None;
    }
    Some(CustomRaftLog {
        data: req.get_custom_request().get_data(),
    })
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
#[repr(u8)]
pub enum CustomRaftLogType {
    Prewrite = 1,
    Commit = 2,
    Rollback = 3,
    PessimisticLock = 4,
    PessimisticRollback = 5,
    OnePc = 6,
    EngineMeta = 7,
    ResolveLock = 8,
    SwitchMemTable = 9,
    TriggerTrimOverBound = 10,
    // 11 is reserved. It was previously TYPE_TRIGGER_MAJOR_COMPACTION.
    TxnFileRef = 12,
    SharedLockWritePut = 13,
    SharedLockLockPut = 14,
    SharedLockLockDelete = 15,
    SharedLockModifies = 16,
    // Note: Please make sure tiflash proxy is updated if you add a new type.
}

impl From<u8> for CustomRaftLogType {
    fn from(v: u8) -> Self {
        match v {
            1 => CustomRaftLogType::Prewrite,
            2 => CustomRaftLogType::Commit,
            3 => CustomRaftLogType::Rollback,
            4 => CustomRaftLogType::PessimisticLock,
            5 => CustomRaftLogType::PessimisticRollback,
            6 => CustomRaftLogType::OnePc,
            7 => CustomRaftLogType::EngineMeta,
            8 => CustomRaftLogType::ResolveLock,
            9 => CustomRaftLogType::SwitchMemTable,
            10 => CustomRaftLogType::TriggerTrimOverBound,
            12 => CustomRaftLogType::TxnFileRef,
            13 => CustomRaftLogType::SharedLockWritePut,
            14 => CustomRaftLogType::SharedLockLockPut,
            15 => CustomRaftLogType::SharedLockLockDelete,
            16 => CustomRaftLogType::SharedLockModifies,
            _ => panic!("unexpected custom raft log type: {:?}", v),
        }
    }
}

const HEADER_SIZE: usize = 2;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
enum SharedLockModifyOp {
    WritePut = 1,
    LockPut = 2,
    LockDelete = 3,
}

impl From<u8> for SharedLockModifyOp {
    fn from(v: u8) -> Self {
        match v {
            1 => SharedLockModifyOp::WritePut,
            2 => SharedLockModifyOp::LockPut,
            3 => SharedLockModifyOp::LockDelete,
            _ => panic!("unexpected shared-lock modify op: {:?}", v),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SharedLockModifyRef<'a> {
    WritePut {
        key: &'a [u8],
        commit_ts: u64,
        value: &'a [u8],
    },
    LockPut {
        key: &'a [u8],
        value: &'a [u8],
    },
    LockDelete {
        key: &'a [u8],
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResolveLockLogEntry<'a> {
    Commit {
        key: &'a [u8],
        commit_ts: u64,
    },
    Rollback {
        key: &'a [u8],
        start_ts: u64,
        delete_lock: bool,
    },
    SharedLockModify(SharedLockModifyRef<'a>),
}

fn decode_key(mut data: &[u8]) -> (&[u8], &[u8]) {
    let key_len = data.get_u16_le() as usize;
    let key = &data[..key_len];
    (&data[key_len..], key)
}

fn decode_value(mut data: &[u8]) -> (&[u8], &[u8]) {
    let val_len = data.get_u32_le() as usize;
    let val = &data[..val_len];
    (&data[val_len..], val)
}

fn decode_lock_put_payload(data: &[u8]) -> (&[u8], SharedLockModifyRef<'_>) {
    let (data, key) = decode_key(data);
    let (data, value) = decode_value(data);
    (data, SharedLockModifyRef::LockPut { key, value })
}

fn decode_lock_delete_payload(data: &[u8]) -> (&[u8], SharedLockModifyRef<'_>) {
    let (data, key) = decode_key(data);
    (data, SharedLockModifyRef::LockDelete { key })
}

fn decode_write_payload(mut data: &[u8]) -> (&[u8], SharedLockModifyRef<'_>) {
    let (rest, key) = decode_key(data);
    data = rest;
    let commit_ts = data.get_u64_le();
    let (data, value) = decode_value(data);
    (
        data,
        SharedLockModifyRef::WritePut {
            key,
            commit_ts,
            value,
        },
    )
}

// CustomRaftLog is the raft log format for unistore to store
// Prewrite/Commit/PessimisticLock.
// | type(1) | reserved(1) | entries
//
// The reserved byte is currently written as zero and is not validated. Rolling
// upgrade safety for new log shapes relies on feature gates preventing new type
// bytes from being emitted before all apply nodes understand them.
//
// It reduces the cost of marshal/unmarshal and avoid DB lookup during apply.
#[derive(Debug)]
pub struct CustomRaftLog<'a> {
    pub(crate) data: &'a [u8],
}

impl<'a> CustomRaftLog<'a> {
    // For debug trace.
    pub fn get_raw(&self) -> &[u8] {
        self.data
    }

    pub fn new_from_data(data: &'a [u8]) -> Self {
        Self { data }
    }

    pub fn get_type(&self) -> CustomRaftLogType {
        self.data[0].into()
    }

    // F: (key, val)
    pub fn iterate_lock<F>(&self, mut f: F)
    where
        F: FnMut(&[u8], &[u8]),
    {
        let mut i = HEADER_SIZE;
        while i < self.data.len() {
            let key_len = LittleEndian::read_u16(&self.data[i..]) as usize;
            i += 2;
            let key = &self.data[i..i + key_len];
            i += key_len;
            let val_len = LittleEndian::read_u32(&self.data[i..]) as usize;
            i += 4;
            let val = &self.data[i..i + val_len];
            i += val_len;
            f(key, val)
        }
    }

    // F: (key, commit_ts)
    pub fn iterate_commit<F>(&self, mut f: F)
    where
        F: FnMut(&[u8], u64),
    {
        let mut i = HEADER_SIZE;
        while i < self.data.len() {
            let key_len = LittleEndian::read_u16(&self.data[i..]) as usize;
            i += 2;
            let key = &self.data[i..i + key_len];
            i += key_len;
            let commit_ts = LittleEndian::read_u64(&self.data[i..]);
            i += 8;
            f(key, commit_ts)
        }
    }

    // F: (key, val, is_extra, del_lock, start_ts, commit_ts)
    pub fn iterate_one_pc<F>(&self, mut f: F)
    where
        F: FnMut(&[u8], &[u8], bool, bool, u64, u64),
    {
        let mut i = HEADER_SIZE;
        while i < self.data.len() {
            let key_len = LittleEndian::read_u16(&self.data[i..]) as usize;
            i += 2;
            let key = &self.data[i..i + key_len];
            i += key_len;
            let val_len = LittleEndian::read_u32(&self.data[i..]) as usize;
            i += 4;
            let val = &self.data[i..i + val_len];
            i += val_len;
            let is_extra = self.data[i] > 0;
            i += 1;
            let del_lock = self.data[i] > 0;
            i += 1;
            let start_ts = LittleEndian::read_u64(&self.data[i..]);
            i += 8;
            let commit_ts = LittleEndian::read_u64(&self.data[i..]);
            i += 8;
            f(key, val, is_extra, del_lock, start_ts, commit_ts)
        }
    }

    // F: (key, start_ts, delete_lock)
    pub fn iterate_rollback<F>(&self, mut f: F)
    where
        F: FnMut(&[u8], u64, bool),
    {
        let mut i = HEADER_SIZE;
        while i < self.data.len() {
            let key_len = LittleEndian::read_u16(&self.data[i..]) as usize;
            i += 2;
            let key = &self.data[i..i + key_len];
            i += key_len;
            let start_ts = LittleEndian::read_u64(&self.data[i..]);
            i += 8;
            let del = self.data[i];
            i += 1;
            f(key, start_ts, del > 0)
        }
    }

    pub fn iterate_del_lock<F>(&self, mut f: F)
    where
        F: FnMut(&[u8]),
    {
        let mut i = HEADER_SIZE;
        while i < self.data.len() {
            let key_len = LittleEndian::read_u16(&self.data[i..]) as usize;
            i += 2;
            let key = &self.data[i..i + key_len];
            i += key_len;
            f(key)
        }
    }

    pub fn iterate_shared_lock_modifies<F>(&self, mut f: F)
    where
        F: FnMut(SharedLockModifyRef<'_>),
    {
        let mut data = &self.data[HEADER_SIZE..];
        while !data.is_empty() {
            let op: SharedLockModifyOp = data.get_u8().into();
            let (rest, modify) = match op {
                SharedLockModifyOp::WritePut => decode_write_payload(data),
                SharedLockModifyOp::LockPut => decode_lock_put_payload(data),
                SharedLockModifyOp::LockDelete => decode_lock_delete_payload(data),
            };
            data = rest;
            f(modify);
        }
    }

    pub(crate) fn get_shared_lock_modify(&self) -> SharedLockModifyRef<'_> {
        let (_, modify) = match self.get_type() {
            CustomRaftLogType::SharedLockWritePut => {
                decode_write_payload(&self.data[HEADER_SIZE..])
            }
            CustomRaftLogType::SharedLockLockPut => {
                decode_lock_put_payload(&self.data[HEADER_SIZE..])
            }
            CustomRaftLogType::SharedLockLockDelete => {
                decode_lock_delete_payload(&self.data[HEADER_SIZE..])
            }
            tp => unreachable!("unexpected shared-lock modify log type: {:?}", tp),
        };
        modify
    }

    pub(crate) fn get_change_set(&self) -> crate::Result<kvenginepb::ChangeSet> {
        let mut cs = kvenginepb::ChangeSet::new();
        cs.merge_from_bytes(&self.data[HEADER_SIZE..])?;
        Ok(cs)
    }

    // The top-level `ResolveLock` log is also reused by cloud_server's
    // `kv_commit` path when commit needs stale pessimistic lock cleanup. In
    // that case, normal commit items are encoded as `Commit`, and delete-only
    // lock cleanup is encoded as `Rollback { start_ts: 0, delete_lock: true }`.
    pub fn iterate_resolve_lock(&self, mut f: impl FnMut(CustomRaftLogType, &[u8], u64, bool)) {
        self.iterate_resolve_lock_entries(|entry| match entry {
            ResolveLockLogEntry::Commit { key, commit_ts } => {
                f(CustomRaftLogType::Commit, key, commit_ts, true)
            }
            ResolveLockLogEntry::Rollback {
                key,
                start_ts,
                delete_lock,
            } => f(CustomRaftLogType::Rollback, key, start_ts, delete_lock),
            ResolveLockLogEntry::SharedLockModify(_) => {
                unreachable!("shared-lock modify entries require iterate_resolve_lock_entries")
            }
        });
    }

    pub fn iterate_resolve_lock_entries<F>(&self, mut f: F)
    where
        F: FnMut(ResolveLockLogEntry<'_>),
    {
        let mut data = &self.data[HEADER_SIZE..];
        while !data.is_empty() {
            let tp: CustomRaftLogType = data.get_u8().into();
            let (rest, entry) = match tp {
                CustomRaftLogType::Commit => {
                    let (data, key) = decode_key(data);
                    let mut data = data;
                    let commit_ts = data.get_u64_le();
                    (data, ResolveLockLogEntry::Commit { key, commit_ts })
                }
                CustomRaftLogType::Rollback => {
                    let (data, key) = decode_key(data);
                    let mut data = data;
                    let start_ts = data.get_u64_le();
                    let delete_lock = data.get_u8() > 0;
                    (
                        data,
                        ResolveLockLogEntry::Rollback {
                            key,
                            start_ts,
                            delete_lock,
                        },
                    )
                }
                CustomRaftLogType::SharedLockWritePut => {
                    let (data, modify) = decode_write_payload(data);
                    (data, ResolveLockLogEntry::SharedLockModify(modify))
                }
                CustomRaftLogType::SharedLockLockPut => {
                    let (data, modify) = decode_lock_put_payload(data);
                    (data, ResolveLockLogEntry::SharedLockModify(modify))
                }
                CustomRaftLogType::SharedLockLockDelete => {
                    let (data, modify) = decode_lock_delete_payload(data);
                    (data, ResolveLockLogEntry::SharedLockModify(modify))
                }
                _ => unreachable!("unexpected custom raft log type: {:?}", tp),
            };
            data = rest;
            f(entry);
        }
    }

    pub(crate) fn get_switch_mem_table(&self) -> u64 {
        let mut bin = &self.data[HEADER_SIZE..];
        bin.get_u64_le()
    }

    pub(crate) fn get_trigger_trim_over_bound(&self) -> TrimOverBoundParameter {
        let bin = &self.data[HEADER_SIZE..];
        TrimOverBoundParameter::unmarshal(bin)
    }

    pub fn is_txn_file_ref(&self) -> bool {
        is_txn_file_ref(self.data)
    }

    pub fn get_txn_file_ref(&self) -> crate::Result<kvenginepb::TxnFileRef> {
        let mut txn_file_ref = kvenginepb::TxnFileRef::new();
        txn_file_ref.merge_from_bytes(&self.data[HEADER_SIZE..])?;
        Ok(txn_file_ref)
    }

    #[allow(clippy::too_many_lines)]
    fn fmt_inner(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self.get_type() {
            CustomRaftLogType::Prewrite => {
                f.write_str("Prewrite { kv_pairs: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_lock(|k, v| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        write!(
                            f,
                            "({}, {})",
                            log_wrappers::Value::key(k),
                            log_wrappers::Value::value(v)
                        )
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::Commit => {
                f.write_str("Commit { items: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_commit(|k, commit_ts| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        write!(
                            f,
                            "Item {{ key: {}, commit_ts: {} }}",
                            log_wrappers::Value::key(k),
                            commit_ts
                        )
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::Rollback => {
                f.write_str("Rollback { items: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_rollback(|k, start_ts, del| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        write!(
                            f,
                            "Item {{ key: {}, start_ts: {}, del_lock: {} }}",
                            log_wrappers::Value::key(k),
                            start_ts,
                            del
                        )
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::PessimisticLock => {
                f.write_str("PessimisticLock { kv_pairs: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_lock(|k, v| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        write!(
                            f,
                            "({}, {})",
                            log_wrappers::Value::key(k),
                            log_wrappers::Value::value(v)
                        )
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::PessimisticRollback => {
                f.write_str("PessimisticRollback { keys: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_del_lock(|k| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        write!(f, "{}", log_wrappers::Value::key(k))
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::OnePc => {
                f.write_str("OnePc { items: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_one_pc(|k, v, is_extra, del_lock, start_ts, commit_ts| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        write!(
                            f,
                            "Item {{ key: {}, value: {}, is_extra: {}, del_lock: {}, start_ts: {}, commit_ts: {} }}",
                            log_wrappers::Value::key(k),
                            log_wrappers::Value::value(v),
                            is_extra,
                            del_lock,
                            start_ts,
                            commit_ts
                        )
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::EngineMeta => f.write_str("EngineMeta"),
            CustomRaftLogType::ResolveLock => {
                f.write_str("ResolveLock { items: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_resolve_lock_entries(|entry| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        match entry {
                            ResolveLockLogEntry::Commit { key, commit_ts } => {
                                Self::fmt_resolve_lock_entry(f, "Commit", key, commit_ts, true)
                            }
                            ResolveLockLogEntry::Rollback {
                                key,
                                start_ts,
                                delete_lock,
                            } => Self::fmt_resolve_lock_entry(
                                f,
                                "Rollback",
                                key,
                                start_ts,
                                delete_lock,
                            ),
                            ResolveLockLogEntry::SharedLockModify(modify) => {
                                Self::fmt_shared_lock_modify(f, modify)
                            }
                        }
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::SharedLockModifies => {
                f.write_str("SharedLockModifies { items: [")?;
                let mut first = true;
                let mut result = Ok(());
                self.iterate_shared_lock_modifies(|modify| {
                    if result.is_err() {
                        return;
                    }
                    result = (|| {
                        if !first {
                            f.write_str(", ")?;
                        }
                        first = false;
                        Self::fmt_shared_lock_modify(f, modify)
                    })();
                });
                result?;
                f.write_str("] }")
            }
            CustomRaftLogType::SharedLockWritePut => {
                let (_, modify) = decode_write_payload(&self.data[HEADER_SIZE..]);
                Self::fmt_shared_lock_modify(f, modify)
            }
            CustomRaftLogType::SharedLockLockPut => {
                let (_, modify) = decode_lock_put_payload(&self.data[HEADER_SIZE..]);
                Self::fmt_shared_lock_modify(f, modify)
            }
            CustomRaftLogType::SharedLockLockDelete => {
                let (_, modify) = decode_lock_delete_payload(&self.data[HEADER_SIZE..]);
                Self::fmt_shared_lock_modify(f, modify)
            }
            CustomRaftLogType::SwitchMemTable => f.write_str("SwitchMemTable"),
            CustomRaftLogType::TriggerTrimOverBound => f.write_str("TriggerTrimOverBound"),
            CustomRaftLogType::TxnFileRef => f.write_str("TxnFileRef"),
        }
    }

    fn fmt_resolve_lock_entry(
        f: &mut Formatter<'_>,
        name: &str,
        key: &[u8],
        ts: u64,
        delete_lock: bool,
    ) -> std::fmt::Result {
        write!(
            f,
            "{} {{ key: {}, ts: {}, del_lock: {} }}",
            name,
            log_wrappers::Value::key(key),
            ts,
            delete_lock
        )
    }

    fn fmt_shared_lock_modify(
        f: &mut Formatter<'_>,
        modify: SharedLockModifyRef<'_>,
    ) -> std::fmt::Result {
        match modify {
            SharedLockModifyRef::WritePut {
                key,
                commit_ts,
                value,
            } => write!(
                f,
                "WritePut {{ key: {}, commit_ts: {}, value: {} }}",
                log_wrappers::Value::key(key),
                commit_ts,
                log_wrappers::Value::value(value)
            ),
            SharedLockModifyRef::LockPut { key, value } => write!(
                f,
                "LockPut {{ key: {}, value: {} }}",
                log_wrappers::Value::key(key),
                log_wrappers::Value::value(value)
            ),
            SharedLockModifyRef::LockDelete { key } => {
                write!(f, "LockDelete {{ key: {} }}", log_wrappers::Value::key(key))
            }
        }
    }

    pub fn is_affect_memtable(
        &self,
        tag: PeerTag,
        base_version: Option<u64>,
        expected_shard_ver: Option<u64>,
    ) -> AffectMemtable {
        match self.get_type() {
            CustomRaftLogType::EngineMeta => {
                let cs = self.get_change_set().unwrap();
                if let Some(expected_shard_ver) = expected_shard_ver {
                    if cs.get_shard_ver() != expected_shard_ver {
                        return AffectMemtable::None;
                    }
                }
                if let Some(snap) = get_any_snap_from_changeset(&cs) {
                    AffectMemtable::Persist {
                        data_seq: snap.data_sequence,
                    }
                } else if cs.has_flush() {
                    AffectMemtable::from_flush(tag, cs.get_flush().version.into(), base_version)
                } else if is_change_set_affect_mem_table(&cs) {
                    AffectMemtable::Write
                } else {
                    AffectMemtable::None
                }
            }
            CustomRaftLogType::SwitchMemTable | CustomRaftLogType::TriggerTrimOverBound => {
                AffectMemtable::None
            }
            _ => AffectMemtable::Write,
        }
    }
}

impl std::fmt::Display for CustomRaftLog<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        self.fmt_inner(f)
    }
}

pub struct CustomBuilder {
    buf: Vec<u8>,
    cnt: i32,
}

impl Default for CustomBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl CustomBuilder {
    pub fn new() -> Self {
        Self {
            buf: vec![0; HEADER_SIZE],
            cnt: 0,
        }
    }

    pub fn append_lock(&mut self, key: &[u8], val: &[u8]) {
        self.buf.put_u16_le(key.len() as u16);
        self.buf.extend_from_slice(key);
        self.buf.put_u32_le(val.len() as u32);
        self.buf.extend_from_slice(val);
        self.cnt += 1;
    }

    pub fn append_commit(&mut self, key: &[u8], commit_ts: u64) {
        self.buf.put_u16_le(key.len() as u16);
        self.buf.extend_from_slice(key);
        self.buf.put_u64_le(commit_ts);
        self.cnt += 1;
    }

    pub fn append_one_pc(
        &mut self,
        key: &[u8],
        val: &[u8],
        is_extra: bool,
        del_lock: bool,
        start_ts: u64,
        commit_ts: u64,
    ) {
        self.buf.put_u16_le(key.len() as u16);
        self.buf.extend_from_slice(key);
        self.buf.put_u32_le(val.len() as u32);
        self.buf.extend_from_slice(val);
        self.buf.put_u8(is_extra as u8);
        self.buf.put_u8(del_lock as u8);
        self.buf.put_u64_le(start_ts);
        self.buf.put_u64_le(commit_ts);
        self.cnt += 1;
    }

    // start_ts == 0 means delete_lock only.
    pub fn append_rollback(&mut self, key: &[u8], start_ts: u64, delete_lock: bool) {
        self.buf.put_u16_le(key.len() as u16);
        self.buf.extend_from_slice(key);
        self.buf.put_u64_le(start_ts);
        self.buf.put_u8(delete_lock as u8);
        self.cnt += 1;
    }

    pub fn append_del_lock(&mut self, key: &[u8]) {
        self.buf.put_u16_le(key.len() as u16);
        self.buf.extend_from_slice(key);
        self.cnt += 1;
    }

    pub fn append_shared_lock_write_put(&mut self, key: &[u8], commit_ts: u64, val: &[u8]) {
        self.buf.put_u16_le(key.len() as u16);
        self.buf.extend_from_slice(key);
        self.buf.put_u64_le(commit_ts);
        self.buf.put_u32_le(val.len() as u32);
        self.buf.extend_from_slice(val);
        self.cnt += 1;
    }

    pub fn append_shared_lock_lock_put(&mut self, key: &[u8], val: &[u8]) {
        self.append_lock(key, val);
    }

    pub fn append_shared_lock_lock_delete(&mut self, key: &[u8]) {
        self.append_del_lock(key);
    }

    pub fn set_change_set(&mut self, cs: &kvenginepb::ChangeSet) {
        assert_eq!(self.buf.len(), HEADER_SIZE);
        let data = cs.write_to_bytes().unwrap();
        self.buf.extend_from_slice(&data);
        self.set_type(CustomRaftLogType::EngineMeta);
    }

    pub fn set_switch_mem_table(&mut self, current_size: u64) {
        assert_eq!(self.buf.len(), HEADER_SIZE);
        self.buf.put_u64_le(current_size);
        self.set_type(CustomRaftLogType::SwitchMemTable);
    }

    pub fn set_trigger_trim_over_bound(&mut self, parameter: &TrimOverBoundParameter) {
        assert_eq!(self.buf.len(), HEADER_SIZE);
        let data = parameter.marshal();
        self.buf.extend_from_slice(&data);
        self.set_type(CustomRaftLogType::TriggerTrimOverBound);
    }

    pub fn set_txn_file(&mut self, txn_file_ref: &kvenginepb::TxnFileRef) {
        assert_eq!(self.buf.len(), HEADER_SIZE);
        let data = txn_file_ref.write_to_bytes().unwrap();
        self.buf.extend_from_slice(&data);
        self.set_type(CustomRaftLogType::TxnFileRef);
    }

    pub fn set_type(&mut self, tp: CustomRaftLogType) {
        self.buf[0] = tp as u8;
    }

    pub fn get_type(&self) -> CustomRaftLogType {
        self.buf[0].into()
    }

    // Some custom logs may contains multiple types of logs, e.g., resolve-lock can
    // contain both commit and rollback. We use type to distinguish them.
    pub fn append_type(&mut self, tp: CustomRaftLogType) {
        self.buf.push(tp as u8);
    }

    fn append_shared_lock_modify_type(&mut self, op: u8) {
        self.buf.push(op);
    }

    pub fn append_shared_lock_modify_write_put(&mut self, key: &[u8], commit_ts: u64, val: &[u8]) {
        self.append_shared_lock_modify_type(SharedLockModifyOp::WritePut as u8);
        self.append_shared_lock_write_put(key, commit_ts, val);
    }

    pub fn append_shared_lock_modify_lock_put(&mut self, key: &[u8], val: &[u8]) {
        self.append_shared_lock_modify_type(SharedLockModifyOp::LockPut as u8);
        self.append_shared_lock_lock_put(key, val);
    }

    pub fn append_shared_lock_modify_lock_delete(&mut self, key: &[u8]) {
        self.append_shared_lock_modify_type(SharedLockModifyOp::LockDelete as u8);
        self.append_shared_lock_lock_delete(key);
    }

    pub fn build(&mut self) -> CustomRequest {
        let mut req = CustomRequest::default();
        let buf = mem::take(&mut self.buf);
        req.set_data(buf);
        req
    }

    pub fn len(&self) -> usize {
        self.cnt as usize
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

pub fn is_engine_meta_log(data: &[u8]) -> bool {
    data[0] == CustomRaftLogType::EngineMeta as u8
}

pub fn is_trigger_trim_over_bound(data: &[u8]) -> bool {
    data[0] == CustomRaftLogType::TriggerTrimOverBound as u8
}

pub fn is_txn_file_ref(data: &[u8]) -> bool {
    data[0] == CustomRaftLogType::TxnFileRef as u8
}

#[derive(Clone, Copy, Default, Debug, PartialEq)]
pub struct TrimOverBoundParameter {
    pub source_shard: Option<IdVer>, // `None` means no trim.
    pub target_shard: Option<IdVer>,
}

impl TrimOverBoundParameter {
    pub fn marshal(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(
            2 + self.source_shard.is_some() as usize * 2 * U64_SIZE
                + self.target_shard.is_some() as usize * 2 * U64_SIZE,
        );

        if let Some(source) = self.source_shard {
            buf.put_u8(1);
            buf.put_u64_le(source.id);
            buf.put_u64_le(source.ver);
        } else {
            buf.put_u8(0);
        }

        if let Some(target) = self.target_shard {
            buf.put_u8(1);
            buf.put_u64_le(target.id);
            buf.put_u64_le(target.ver);
        } else {
            buf.put_u8(0);
        }

        buf
    }

    pub fn unmarshal(mut data: &[u8]) -> Self {
        if data.is_empty() {
            return Self::default();
        }
        let source_shard = if data.get_u8() == 1 {
            Some(IdVer::new(data.get_u64_le(), data.get_u64_le()))
        } else {
            None
        };
        let target_shard = if data.get_u8() == 1 {
            Some(IdVer::new(data.get_u64_le(), data.get_u64_le()))
        } else {
            None
        };
        Self {
            source_shard,
            target_shard,
        }
    }

    pub fn source_shard_id(&self) -> Option<u64> {
        self.source_shard.map(|x| x.id)
    }

    pub fn target_shard_id(&self) -> Option<u64> {
        self.target_shard.map(|x| x.id)
    }

    pub fn is_for_shard(&self, shard_id: u64) -> bool {
        Some(shard_id) == self.source_shard_id() || Some(shard_id) == self.target_shard_id()
    }
}

#[derive(Debug)]
pub enum AffectMemtable {
    None,
    Write,
    Persist { data_seq: u64 },
}

impl AffectMemtable {
    fn from_flush(tag: PeerTag, snap_version: SnapVersion, base_version: Option<u64>) -> Self {
        let Some(base_version) = base_version else {
            // Shard is not initialized yet. Should not happen.
            warn!("{} AffectMemtable: no base version", tag);
            debug_assert!(false);
            return Self::None;
        };
        let Some(data_seq) = snap_version.into_inner().checked_sub(base_version) else {
            warn!("{} AffectMemtable: invalid snap version", tag;
                "snap_ver" => snap_version, "base_ver" => base_version);
            debug_assert!(false);
            return Self::None;
        };
        Self::Persist { data_seq }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_custom_log() {
        let mut builder = CustomBuilder::new();
        builder.set_switch_mem_table(2022);
        let req = builder.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(cl.get_type(), CustomRaftLogType::SwitchMemTable);
        assert_eq!(cl.get_switch_mem_table(), 2022);
    }

    #[test]
    fn test_trim_over_bound_parameter() {
        let cases = vec![
            (None, None),
            (Some(IdVer { id: 1, ver: 2 }), None),
            (None, Some(IdVer { id: 2, ver: 3 })),
            (Some(IdVer { id: 1, ver: 2 }), Some(IdVer { id: 2, ver: 3 })),
        ];

        for (source_shard, target_shard) in cases {
            let trim_over_bound = TrimOverBoundParameter {
                source_shard,
                target_shard,
            };
            assert_eq!(
                trim_over_bound,
                TrimOverBoundParameter::unmarshal(trim_over_bound.marshal().as_slice())
            );
        }

        assert_eq!(
            TrimOverBoundParameter::default(),
            TrimOverBoundParameter::unmarshal(vec![].as_slice())
        );
    }

    #[test]
    fn test_custom_raft_log_types() {
        let valid_types = [
            (1, CustomRaftLogType::Prewrite),
            (2, CustomRaftLogType::Commit),
            (3, CustomRaftLogType::Rollback),
            (4, CustomRaftLogType::PessimisticLock),
            (5, CustomRaftLogType::PessimisticRollback),
            (6, CustomRaftLogType::OnePc),
            (7, CustomRaftLogType::EngineMeta),
            (8, CustomRaftLogType::ResolveLock),
            (9, CustomRaftLogType::SwitchMemTable),
            (10, CustomRaftLogType::TriggerTrimOverBound),
            (12, CustomRaftLogType::TxnFileRef),
            (13, CustomRaftLogType::SharedLockWritePut),
            (14, CustomRaftLogType::SharedLockLockPut),
            (15, CustomRaftLogType::SharedLockLockDelete),
            (16, CustomRaftLogType::SharedLockModifies),
        ];
        for (i, expected) in valid_types {
            let tp: CustomRaftLogType = i.into();
            assert_eq!(expected, tp);
            assert_eq!(i, tp as u8);
        }

        // Test invalid and reserved types will panic.
        for i in [0, 11, 17, 18, 19, 20] {
            let result = std::panic::catch_unwind(|| {
                let _: CustomRaftLogType = i.into();
            });
            assert!(result.is_err());
        }
    }

    #[test]
    fn test_shared_lock_modifies_log() {
        #[derive(Debug, PartialEq, Eq)]
        enum OwnedSharedLockModify {
            WritePut {
                key: Vec<u8>,
                commit_ts: u64,
                value: Vec<u8>,
            },
            LockPut {
                key: Vec<u8>,
                value: Vec<u8>,
            },
            LockDelete {
                key: Vec<u8>,
            },
        }

        let mut builder = CustomBuilder::new();
        builder.set_type(CustomRaftLogType::SharedLockModifies);
        builder.append_shared_lock_modify_write_put(b"k1", 42, b"write");
        builder.append_shared_lock_modify_lock_put(b"k2", b"lock");
        builder.append_shared_lock_modify_lock_delete(b"k3");

        let req = builder.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(cl.get_type(), CustomRaftLogType::SharedLockModifies);

        let mut modifies = vec![];
        cl.iterate_shared_lock_modifies(|modify| {
            modifies.push(match modify {
                SharedLockModifyRef::WritePut {
                    key,
                    commit_ts,
                    value,
                } => OwnedSharedLockModify::WritePut {
                    key: key.to_vec(),
                    commit_ts,
                    value: value.to_vec(),
                },
                SharedLockModifyRef::LockPut { key, value } => OwnedSharedLockModify::LockPut {
                    key: key.to_vec(),
                    value: value.to_vec(),
                },
                SharedLockModifyRef::LockDelete { key } => {
                    OwnedSharedLockModify::LockDelete { key: key.to_vec() }
                }
            });
        });
        assert_eq!(
            modifies,
            vec![
                OwnedSharedLockModify::WritePut {
                    key: b"k1".to_vec(),
                    commit_ts: 42,
                    value: b"write".to_vec(),
                },
                OwnedSharedLockModify::LockPut {
                    key: b"k2".to_vec(),
                    value: b"lock".to_vec(),
                },
                OwnedSharedLockModify::LockDelete {
                    key: b"k3".to_vec(),
                },
            ]
        );
    }

    #[test]
    fn test_is_affect_memtable_with_expected_shard_ver() {
        let mut cs = kvenginepb::ChangeSet::new();
        cs.set_shard_id(7);
        cs.set_shard_ver(11);
        cs.mut_flush().set_version(120);
        let mut builder = CustomBuilder::new();
        builder.set_change_set(&cs);
        let req = builder.build();
        let custom_log = CustomRaftLog::new_from_data(req.get_data());
        let tag = PeerTag::default();

        assert!(matches!(
            custom_log.is_affect_memtable(tag, Some(100), None),
            AffectMemtable::Persist { data_seq: 20 }
        ));
        assert!(matches!(
            custom_log.is_affect_memtable(tag, Some(100), Some(11)),
            AffectMemtable::Persist { data_seq: 20 }
        ));
        assert!(matches!(
            custom_log.is_affect_memtable(tag, Some(100), Some(12)),
            AffectMemtable::None
        ));
    }

    #[test]
    fn test_custom_raft_log_display() {
        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::Prewrite);
        b.append_lock(b"k1", b"v1");
        b.append_lock(b"k2", b"v2");
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "Prewrite { kv_pairs: [(6B31, 7631), (6B32, 7632)] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::Commit);
        b.append_commit(b"k3", 100);
        b.append_commit(b"k4", 100);
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "Commit { items: [Item { key: 6B33, commit_ts: 100 }, \
            Item { key: 6B34, commit_ts: 100 }] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::Rollback);
        b.append_rollback(b"k5", 110, true);
        b.append_rollback(b"k6", 110, false);
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "Rollback { items: [Item { key: 6B35, start_ts: 110, del_lock: true }, \
            Item { key: 6B36, start_ts: 110, del_lock: false }] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::PessimisticLock);
        b.append_lock(b"k7", b"v7");
        b.append_lock(b"k8", b"v8");
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "PessimisticLock { kv_pairs: [(6B37, 7637), (6B38, 7638)] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::PessimisticRollback);
        b.append_del_lock(b"k9");
        b.append_del_lock(b"k10");
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "PessimisticRollback { keys: [6B39, 6B3130] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::OnePc);
        b.append_one_pc(b"k11", b"v11", false, true, 120, 130);
        b.append_one_pc(b"k12", b"v12", true, false, 120, 130);
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "OnePc { items: [Item { key: 6B3131, value: 763131, is_extra: false, del_lock: true, start_ts: 120, commit_ts: 130 }, \
            Item { key: 6B3132, value: 763132, is_extra: true, del_lock: false, start_ts: 120, commit_ts: 130 }] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::ResolveLock);
        b.append_type(CustomRaftLogType::Commit);
        b.append_commit(b"k13", 140);
        b.append_type(CustomRaftLogType::Rollback);
        b.append_rollback(b"k14", 150, false);
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "ResolveLock { items: [Commit { key: 6B3133, ts: 140, del_lock: true }, \
            Rollback { key: 6B3134, ts: 150, del_lock: false }] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::SharedLockModifies);
        b.append_shared_lock_modify_write_put(b"k15", 160, b"v15");
        b.append_shared_lock_modify_lock_put(b"k16", b"v16");
        b.append_shared_lock_modify_lock_delete(b"k17");
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "SharedLockModifies { items: [WritePut { key: 6B3135, commit_ts: 160, value: 763135 }, \
            LockPut { key: 6B3136, value: 763136 }, LockDelete { key: 6B3137 }] }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::SharedLockWritePut);
        b.append_shared_lock_write_put(b"k18", 170, b"v18");
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(
            format!("{}", cl),
            "WritePut { key: 6B3138, commit_ts: 170, value: 763138 }"
        );

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::SharedLockLockPut);
        b.append_shared_lock_lock_put(b"k19", b"v19");
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(format!("{}", cl), "LockPut { key: 6B3139, value: 763139 }");

        let mut b = CustomBuilder::new();
        b.set_type(CustomRaftLogType::SharedLockLockDelete);
        b.append_shared_lock_lock_delete(b"k20");
        let req = b.build();
        let cl = CustomRaftLog::new_from_data(req.get_data());
        assert_eq!(format!("{}", cl), "LockDelete { key: 6B3230 }");
    }
}