// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
use std::mem;

use tikv_util::{Either, box_err};
use txn_types::{Key, LockType, ReqType, TimeStamp};

use crate::storage::{
    ProcessResult, Result as StorageResult, Snapshot,
    kv::WriteData,
    lock_manager::LockManager,
    mvcc::{
        ErrorInner as MvccErrorInner, MvccTxn, Result as MvccResult, SharedLockMutation,
        SnapshotReader,
    },
    txn::{
        Result,
        commands::{
            Command, CommandExt, PessimisticRollbackReadPhase, ReaderWithStats, ReleasedLocks,
            ResponsePolicy, TypedCommand, WriteCommand, WriteContext, WriteResult,
        },
    },
};

command! {
    /// Rollback pessimistic locks identified by `start_ts` and `for_update_ts`.
    ///
    /// This can roll back an [`AcquirePessimisticLock`](Command::AcquirePessimisticLock) command.
    PessimisticRollback:
        cmd_ty => Vec<StorageResult<()>>,
        display => "kv::command::pessimistic_rollback keys({:?}) @ {} {} | {:?}", (keys, start_ts, for_update_ts, ctx),
        content => {
            /// The keys to be rolled back.
            keys: Vec<Key>,
            /// The transaction timestamp.
            start_ts: TimeStamp,
            for_update_ts: TimeStamp,
            /// The next key to scan using pessimistic rollback read phase.
            scan_key: Option<Key>,
        }
}

impl CommandExt for PessimisticRollback {
    ctx!();
    tag!(pessimistic_rollback);
    request_type!(KvPessimisticRollback);
    ts!(start_ts);
    write_bytes!(keys: multiple);
    gen_lock!(keys: multiple);
}

#[maybe_async::async_trait]
impl<S: Snapshot + 'static, L: LockManager> WriteCommand<S, L> for PessimisticRollback {
    /// Delete any pessimistic lock with small for_update_ts belongs to this
    /// transaction.
    #[maybe_async]
    async fn process_write(
        mut self,
        snapshot: S,
        context: WriteContext<'_, L>,
    ) -> Result<WriteResult> {
        let mut txn = MvccTxn::new(self.start_ts, context.concurrency_manager);
        let mut reader = ReaderWithStats::new(
            SnapshotReader::new_with_ctx(self.start_ts, snapshot, &self.ctx),
            context.statistics,
        );

        let ctx = mem::take(&mut self.ctx);
        let keys = mem::take(&mut self.keys);

        let rows = keys.len();
        let mut released_locks = ReleasedLocks::new();
        for key in keys {
            fail_point!("pessimistic_rollback", |err| Err(
                crate::storage::mvcc::Error::from(crate::storage::mvcc::txn::make_txn_error(
                    err,
                    &key,
                    self.start_ts
                ))
                .into()
            ));
            let released_lock: MvccResult<_> = if let Some(lock_or_shared) =
                reader.load_lock(&key)?
            {
                match lock_or_shared {
                    Either::Left(lock)
                        if lock.lock_type == LockType::Pessimistic
                            && lock.ts == self.start_ts
                            && lock.for_update_ts <= self.for_update_ts =>
                    {
                        Ok(txn.unlock_key(key, true, TimeStamp::zero()))
                    }
                    Either::Left(_) => Ok(None),
                    Either::Right(mut shared_locks) => {
                        txn.mark_shared_lock_lifecycle();
                        let should_rollback = match shared_locks
                            .get_lock(&self.start_ts)
                            .map_err(crate::storage::mvcc::Error::from)?
                        {
                            Some(lock) if key.is_encoded_from(&lock.primary) => {
                                return Err(crate::storage::mvcc::Error::from(
                                        MvccErrorInner::Other(box_err!(
                                            "shared pessimistic lock must not be used as transaction primary"
                                        )),
                                    )
                                    .into());
                            }
                            Some(lock) => {
                                lock.lock_type == LockType::Pessimistic
                                    && lock.for_update_ts <= self.for_update_ts
                            }
                            None => false,
                        };

                        if should_rollback {
                            shared_locks
                                .remove_lock(&self.start_ts)
                                .map_err(crate::storage::mvcc::Error::from)?;
                            if shared_locks.is_empty() {
                                Ok(txn.delete_shared_lock_wrapper(key, TimeStamp::zero()))
                            } else {
                                txn.put_shared_locks(
                                    key,
                                    shared_locks,
                                    SharedLockMutation::HolderRemoved,
                                );
                                Ok(None)
                            }
                        } else {
                            Ok(None)
                        }
                    }
                }
            } else {
                Ok(None)
            };
            released_locks.push(released_lock?);
        }

        let pr = if self.scan_key.is_none() {
            ProcessResult::MultiRes { results: vec![] }
        } else {
            let next_cmd = PessimisticRollbackReadPhase {
                ctx: ctx.clone(),
                deadline: self.deadline,
                start_ts: self.start_ts,
                for_update_ts: self.for_update_ts,
                scan_key: self.scan_key.take(),
            };
            ProcessResult::NextCommand {
                cmd: Command::PessimisticRollbackReadPhase(next_cmd),
            }
        };

        let shared_lock_upgrade_hints = txn.take_shared_lock_upgrade_hints();
        let shared_lock_lifecycle = txn.has_shared_lock_lifecycle();
        let mut write_data = WriteData::from_modifies(txn.into_modifies());
        write_data.set_req_type(ReqType::PessimisticRollback);
        write_data.set_shared_lock_lifecycle(shared_lock_lifecycle);
        write_data.set_allowed_on_disk_almost_full();
        Ok(WriteResult {
            ctx,
            to_be_write: write_data,
            rows,
            pr,
            lock_info: vec![],
            released_locks,
            new_acquired_locks: vec![],
            shared_lock_upgrade_hints,
            lock_guards: vec![],
            response_policy: ResponsePolicy::OnApplied,
            known_txn_status: vec![],
        })
    }
}

#[cfg(test)]
pub mod tests {
    use cloud_server::mock_kv_engine::TestKvEngine;
    use concurrency_manager::ConcurrencyManager;
    use kvproto::kvrpcpb::Context;
    use tikv_util::deadline::Deadline;
    use txn_types::Key;

    use super::*;
    use crate::storage::{
        TestEngineBuilder,
        kv::Engine,
        lock_manager::MockLockManager,
        mvcc::tests::*,
        txn::{
            commands::{WriteCommand, WriteContext},
            scheduler::DEFAULT_EXECUTION_DURATION_LIMIT,
            tests::*,
            txn_status_cache::TxnStatusCache,
        },
    };

    #[cfg(test)]
    fn write_shared_pessimistic_locks<E: Engine>(
        engine: &E,
        key: &[u8],
        primary: &[u8],
        holders: &[(u64, u64)],
    ) {
        let mut shared_locks = txn_types::SharedLocks::new();
        for (start_ts, for_update_ts) in holders {
            shared_locks
                .insert_lock(txn_types::Lock::new(
                    LockType::Pessimistic,
                    primary.to_vec(),
                    (*start_ts).into(),
                    3000,
                    None,
                    (*for_update_ts).into(),
                    1,
                    TimeStamp::zero(),
                ))
                .unwrap();
        }
        let mut data = tikv_kv::WriteData::from_modifies(vec![tikv_kv::Modify::Put(
            engine_traits::CF_LOCK,
            Key::from_raw(key),
            shared_locks.to_bytes(),
        )]);
        data.extra.req_type = txn_types::ReqType::PessimisticLock;
        data.set_shared_lock_lifecycle(true);
        write(engine, &Context::default(), data);
    }

    #[cfg(test)]
    fn must_shared_locks<E: Engine>(engine: &mut E, key: &[u8]) -> txn_types::SharedLocks {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let mut reader = SnapshotReader::new(TimeStamp::zero(), snapshot, true);
        match reader.load_lock(&Key::from_raw(key)).unwrap().unwrap() {
            Either::Right(shared_locks) => shared_locks,
            Either::Left(lock) => panic!("expected shared locks, got {:?}", lock),
        }
    }

    #[cfg(test)]
    fn must_unlocked_from_cloud(engine: &mut TestKvEngine, key: &[u8]) {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let mut reader = SnapshotReader::new(TimeStamp::zero(), snapshot, true);
        assert!(reader.load_lock(&Key::from_raw(key)).unwrap().is_none());
    }

    pub fn must_success<E: Engine>(
        engine: &mut E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        for_update_ts: impl Into<TimeStamp>,
    ) {
        let ctx = Context::default();
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let for_update_ts = for_update_ts.into();
        let cm = ConcurrencyManager::new(for_update_ts);
        let start_ts = start_ts.into();
        let command = crate::storage::txn::commands::PessimisticRollback {
            ctx: ctx.clone(),
            keys: vec![Key::from_raw(key)],
            start_ts,
            for_update_ts,
            deadline: Deadline::from_now(DEFAULT_EXECUTION_DURATION_LIMIT),
            scan_key: None,
        };
        let lock_mgr = MockLockManager::new();
        let write_context = WriteContext {
            lock_mgr: &lock_mgr,
            concurrency_manager: cm,
            extra_op: Default::default(),
            extra_cf_scan_mode: Default::default(),
            statistics: &mut Default::default(),
            async_apply_prewrite: false,
            txn_status_cache: &TxnStatusCache::new_for_test(),
        };
        let result = command.process_write(snapshot, write_context).unwrap();
        let shared_lock_lifecycle = result.to_be_write.extra.shared_lock_lifecycle;
        let mut data = tikv_kv::WriteData::from_modifies(result.to_be_write.modifies);
        data.extra.req_type = txn_types::ReqType::PessimisticRollback;
        data.set_shared_lock_lifecycle(shared_lock_lifecycle);
        write(engine, &ctx, data);
    }

    pub fn must_err<E: Engine>(
        engine: &mut E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        for_update_ts: impl Into<TimeStamp>,
    ) -> crate::storage::txn::Error {
        let ctx = Context::default();
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let for_update_ts = for_update_ts.into();
        let cm = ConcurrencyManager::new(for_update_ts);
        let start_ts = start_ts.into();
        let command = crate::storage::txn::commands::PessimisticRollback {
            ctx: ctx.clone(),
            keys: vec![Key::from_raw(key)],
            start_ts,
            for_update_ts,
            deadline: Deadline::from_now(DEFAULT_EXECUTION_DURATION_LIMIT),
            scan_key: None,
        };
        let lock_mgr = MockLockManager::new();
        let write_context = WriteContext {
            lock_mgr: &lock_mgr,
            concurrency_manager: cm,
            extra_op: Default::default(),
            extra_cf_scan_mode: Default::default(),
            statistics: &mut Default::default(),
            async_apply_prewrite: false,
            txn_status_cache: &TxnStatusCache::new_for_test(),
        };
        match command.process_write(snapshot, write_context) {
            Ok(_) => panic!("expected pessimistic rollback to fail"),
            Err(err) => err,
        }
    }

    #[test]
    fn test_pessimistic_rollback_shared_lock_emits_upgrade_hint_for_remaining_holder() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"pessimistic-rollback-upgrade-hint";
        write_shared_pessimistic_locks(&engine, key, b"primary", &[(10, 10), (20, 20)]);
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let lock_mgr = MockLockManager::new();
        let command = crate::storage::txn::commands::PessimisticRollback {
            ctx: Context::default(),
            keys: vec![Key::from_raw(key)],
            start_ts: 20.into(),
            for_update_ts: 20.into(),
            deadline: Deadline::from_now(DEFAULT_EXECUTION_DURATION_LIMIT),
            scan_key: None,
        };
        let result = command
            .process_write(
                snapshot,
                WriteContext {
                    lock_mgr: &lock_mgr,
                    concurrency_manager: ConcurrencyManager::new(TimeStamp::zero()),
                    extra_op: Default::default(),
                    extra_cf_scan_mode: Default::default(),
                    statistics: &mut Default::default(),
                    async_apply_prewrite: false,
                    txn_status_cache: &TxnStatusCache::new_for_test(),
                },
            )
            .unwrap();

        assert_eq!(result.shared_lock_upgrade_hints.len(), 1);
        assert_eq!(result.shared_lock_upgrade_hints[0].key, Key::from_raw(key));
        assert_eq!(
            result.shared_lock_upgrade_hints[0].sole_holder_start_ts,
            10.into()
        );
    }

    #[test]
    fn test_pessimistic_rollback() {
        let mut engine = TestEngineBuilder::new().build().unwrap();

        let k = b"k1";
        let v = b"v1";

        // Normal
        must_acquire_pessimistic_lock(&mut engine, k, k, 1, 1);
        must_pessimistic_locked(&mut engine, k, 1, 1);
        must_success(&mut engine, k, 1, 1);
        must_unlocked(&mut engine, k);
        must_get_commit_ts_none(&mut engine, k, 1);
        // Pessimistic rollback is idempotent
        must_success(&mut engine, k, 1, 1);
        must_unlocked(&mut engine, k);
        must_get_commit_ts_none(&mut engine, k, 1);

        // Succeed if the lock doesn't exist.
        must_success(&mut engine, k, 2, 2);

        // Do nothing if meets other transaction's pessimistic lock
        must_acquire_pessimistic_lock(&mut engine, k, k, 2, 3);
        must_success(&mut engine, k, 1, 1);
        must_success(&mut engine, k, 1, 2);
        must_success(&mut engine, k, 1, 3);
        must_success(&mut engine, k, 1, 4);
        must_success(&mut engine, k, 3, 3);
        must_success(&mut engine, k, 4, 4);

        // Succeed if for_update_ts is larger; do nothing if for_update_ts is smaller.
        must_pessimistic_locked(&mut engine, k, 2, 3);
        must_success(&mut engine, k, 2, 2);
        must_pessimistic_locked(&mut engine, k, 2, 3);
        must_success(&mut engine, k, 2, 4);
        must_unlocked(&mut engine, k);

        // Do nothing if rollbacks a non-pessimistic lock.
        must_prewrite_put(&mut engine, k, v, k, 3);
        must_locked(&mut engine, k, 3);
        must_success(&mut engine, k, 3, 3);
        must_locked(&mut engine, k, 3);

        // Do nothing if meets other transaction's optimistic lock
        must_success(&mut engine, k, 2, 2);
        must_success(&mut engine, k, 2, 3);
        must_success(&mut engine, k, 2, 4);
        must_success(&mut engine, k, 4, 4);
        must_locked(&mut engine, k, 3);

        // Do nothing if committed
        must_commit(&mut engine, k, 3, 4);
        must_unlocked(&mut engine, k);
        must_get_commit_ts(&mut engine, k, 3, 4);
        must_success(&mut engine, k, 3, 3);
        must_success(&mut engine, k, 3, 4);
        must_success(&mut engine, k, 3, 5);
    }

    #[test]
    fn test_pessimistic_rollback_shared_lock_removes_only_matching_holder() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-pessimistic-rollback";
        let primary = b"primary";
        write_shared_pessimistic_locks(&engine, key, primary, &[(10, 15), (20, 25)]);

        must_success(&mut engine, key, 10, 14);
        let shared_locks = must_shared_locks(&mut engine, key);
        assert!(shared_locks.contains_start_ts(10.into()));
        assert!(shared_locks.contains_start_ts(20.into()));

        must_success(&mut engine, key, 10, 15);
        let shared_locks = must_shared_locks(&mut engine, key);
        assert!(!shared_locks.contains_start_ts(10.into()));
        assert!(shared_locks.contains_start_ts(20.into()));
        must_get_commit_ts_none(&mut engine, key, 10);

        must_success(&mut engine, key, 20, 25);
        must_unlocked_from_cloud(&mut engine, key);
        must_get_commit_ts_none(&mut engine, key, 20);
    }

    #[test]
    fn test_pessimistic_rollback_rejects_shared_lock_primary() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-pessimistic-rollback-primary";
        write_shared_pessimistic_locks(&engine, key, key, &[(10, 15)]);

        let err = must_err(&mut engine, key, 10, 15);
        match err {
            crate::storage::txn::Error(box crate::storage::txn::ErrorInner::Mvcc(
                crate::storage::mvcc::Error(box MvccErrorInner::Other(err)),
            )) => {
                assert!(
                    err.to_string().contains(
                        "shared pessimistic lock must not be used as transaction primary"
                    )
                );
            }
            err => panic!("unexpected error: {:?}", err),
        }
    }
}