// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
use tikv_util::Either;

use crate::storage::{
    Snapshot, TxnStatus,
    mvcc::{
        ErrorInner, Key, MvccTxn, ReleasedLock, Result as MvccResult, SnapshotReader, TimeStamp,
        metrics::{MVCC_CONFLICT_COUNTER, MVCC_DUPLICATE_CMD_COUNTER_VEC},
    },
    txn::actions::check_txn_status::{
        MissingLockAction, check_txn_status_missing_lock, check_txn_status_missing_lock_async,
        rollback_lock, rollback_lock_async, rollback_shared_lock, rollback_shared_lock_async,
    },
};

/// Cleanup the lock if it's TTL has expired, comparing with `current_ts`. If
/// `current_ts` is 0, cleanup the lock without checking TTL. If the lock is the
/// primary lock of a pessimistic transaction, the rollback record is protected
/// from being collapsed.
///
/// Returns the released lock. Returns error if the key is locked or has already
/// been committed.
#[maybe_async::both]
pub async fn cleanup<S: Snapshot>(
    txn: &mut MvccTxn,
    reader: &mut SnapshotReader<S>,
    key: Key,
    current_ts: TimeStamp,
    protect_rollback: bool,
) -> MvccResult<Option<ReleasedLock>> {
    fail_point!("cleanup", |err| Err(
        crate::storage::mvcc::txn::make_txn_error(err, &key, reader.start_ts).into()
    ));

    let lock_state = match txn.get_pending_lock_bytes(&key) {
        Some(None) => None,
        Some(Some(bytes)) => Some(txn_types::parse_lock(bytes)?),
        None => reader.load_lock(&key)?,
    };

    match lock_state {
        Some(Either::Left(ref lock)) if lock.ts == reader.start_ts => {
            // If current_ts is not 0, check the Lock's TTL.
            // If the lock is not expired, do not rollback it but report key is locked.
            // If current lock is from txn file, we also reject normal resolve_lock request.
            if lock.is_txn_file
                || (!current_ts.is_zero() && lock.ts.physical() + lock.ttl >= current_ts.physical())
            {
                return Err(
                    ErrorInner::KeyIsLocked(lock.clone().into_lock_info(key.into_raw()?)).into(),
                );
            }
            rollback_lock(
                txn,
                reader,
                key,
                lock,
                lock.is_pessimistic_txn(),
                !protect_rollback,
            )
            .await
        }
        Some(Either::Right(mut shared_locks))
            if shared_locks.contains_start_ts(reader.start_ts) =>
        {
            if !current_ts.is_zero() {
                let lock = shared_locks
                    .get_lock(&reader.start_ts)
                    .map_err(crate::storage::mvcc::Error::from)?
                    .expect("shared lock entry must exist for current txn");
                let expire_time = lock.ts.physical() + lock.ttl;
                if expire_time >= current_ts.physical() {
                    return Err(ErrorInner::KeyIsLocked(
                        shared_locks.into_lock_info(key.into_raw()?),
                    )
                    .into());
                }
            }
            rollback_shared_lock(
                txn,
                reader,
                key,
                shared_locks,
                reader.start_ts,
                !protect_rollback,
            )
            .await
        }
        l => match check_txn_status_missing_lock(
            txn,
            reader,
            key.clone(),
            l.and_then(|lock_or_shared| lock_or_shared.left()),
            MissingLockAction::rollback_protect(protect_rollback),
            false,
        )
        .await?
        {
            TxnStatus::Committed { commit_ts } => {
                MVCC_CONFLICT_COUNTER.rollback_committed.inc();
                Err(ErrorInner::Committed {
                    start_ts: reader.start_ts,
                    commit_ts,
                    key: key.into_raw()?,
                }
                .into())
            }
            TxnStatus::RolledBack => {
                // Return Ok on Rollback already exist.
                MVCC_DUPLICATE_CMD_COUNTER_VEC.rollback.inc();
                Ok(None)
            }
            TxnStatus::LockNotExist => Ok(None),
            _ => unreachable!(),
        },
    }
}

pub mod tests {
    #[cfg(test)]
    use cloud_server::mock_kv_engine::TestKvEngine;
    use concurrency_manager::ConcurrencyManager;
    use engine_traits::CF_WRITE;
    use kvproto::kvrpcpb::Context;
    #[cfg(test)]
    use kvproto::kvrpcpb::PrewriteRequestPessimisticAction::*;
    use tikv_kv::WriteData;
    use txn_types::{ReqType, TimeStamp};

    use super::*;
    use crate::storage::{
        Engine,
        mvcc::{
            Error as MvccError, WriteType,
            tests::{must_have_write, must_not_have_write, write},
        },
        txn::tests::{must_commit, must_prewrite_put},
    };
    #[cfg(test)]
    use crate::storage::{
        TestEngineBuilder,
        mvcc::tests::{
            must_get_rollback_protected, must_get_rollback_ts, must_locked, must_unlocked,
            must_written,
        },
        txn::commands::txn_heart_beat,
        txn::tests::{must_acquire_pessimistic_lock, must_pessimistic_prewrite_put},
    };

    #[cfg(test)]
    fn write_shared_prewrite_locks<E: Engine>(
        engine: &E,
        key: &[u8],
        primary: &[u8],
        holders: &[u64],
    ) {
        write_shared_prewrite_locks_with_flags(engine, key, primary, holders, false);
    }

    #[cfg(test)]
    fn write_shared_prewrite_locks_with_flags<E: Engine>(
        engine: &E,
        key: &[u8],
        primary: &[u8],
        holders: &[u64],
        use_async_commit: bool,
    ) {
        let mut shared_locks = txn_types::SharedLocks::new();
        for start_ts in holders {
            let mut lock = txn_types::Lock::new(
                txn_types::LockType::Lock,
                primary.to_vec(),
                (*start_ts).into(),
                3000,
                None,
                TimeStamp::zero(),
                1,
                TimeStamp::zero(),
            );
            lock.use_async_commit = use_async_commit;
            if use_async_commit {
                lock.secondaries = vec![b"secondary".to_vec()];
            }
            shared_locks.insert_lock(lock).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 = 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_succeed<E: Engine>(
        engine: &mut E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        current_ts: impl Into<TimeStamp>,
    ) {
        let ctx = Context::default();
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let current_ts = current_ts.into();
        let cm = ConcurrencyManager::new(current_ts);
        let start_ts = start_ts.into();
        let mut txn = MvccTxn::new(start_ts, cm);
        let mut reader = SnapshotReader::new(start_ts, snapshot, true);
        cleanup(&mut txn, &mut reader, Key::from_raw(key), current_ts, true).unwrap();
        let shared_lock_lifecycle = txn.has_shared_lock_lifecycle();
        let mut data = WriteData::from_modifies(txn.into_modifies());
        // TODO: determine this type.
        data.extra.req_type = ReqType::ResolveLock;
        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>,
        current_ts: impl Into<TimeStamp>,
    ) -> MvccError {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let current_ts = current_ts.into();
        let cm = ConcurrencyManager::new(current_ts);
        let start_ts = start_ts.into();
        let mut txn = MvccTxn::new(start_ts, cm);
        let mut reader = SnapshotReader::new(start_ts, snapshot, true);
        cleanup(&mut txn, &mut reader, Key::from_raw(key), current_ts, true).unwrap_err()
    }

    pub fn must_cleanup_with_gc_fence<E: Engine>(
        engine: &mut E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        current_ts: impl Into<TimeStamp>,
        gc_fence: impl Into<TimeStamp>,
        without_target_write: bool,
    ) {
        let ctx = Context::default();
        let gc_fence = gc_fence.into();
        let start_ts = start_ts.into();
        let current_ts = current_ts.into();

        if !gc_fence.is_zero() && without_target_write {
            // NOTE: kvengine does not support delete key directly with the original commit
            // ts, so the gc fence handle need 3 ts: prewrite,
            // commit(prewrite+1),delete(commit+1). Put a dummy record and
            // remove it after doing cleanup.
            must_not_have_write(engine, key, gc_fence);
            must_prewrite_put(engine, key, b"dummy_value", key, gc_fence.prev());
            must_commit(engine, key, gc_fence.prev(), gc_fence);
        }

        let cm = ConcurrencyManager::new(current_ts);
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let mut txn = MvccTxn::new(start_ts, cm);
        let mut reader = SnapshotReader::new(start_ts, snapshot, true);
        cleanup(&mut txn, &mut reader, Key::from_raw(key), current_ts, true).unwrap();

        let mut data = tikv_kv::WriteData::from_modifies(txn.into_modifies());
        // TODO: determine the correct req type.
        data.extra.req_type = txn_types::ReqType::Rollback;
        write(engine, &ctx, data);

        let w = must_have_write(engine, key, start_ts);
        assert_ne!(w.start_ts, start_ts, "no overlapping write record");
        assert!(
            w.write_type != WriteType::Rollback && w.write_type != WriteType::Lock,
            "unexpected write type {:?}",
            w.write_type
        );

        if !gc_fence.is_zero() && without_target_write {
            engine
                .delete_cf(&ctx, CF_WRITE, Key::from_raw(key).append_ts(gc_fence))
                .unwrap();
            must_not_have_write(engine, key, gc_fence);
        }
    }

    #[test]
    #[ignore = "CSE incompatibility; see #6158"]
    fn test_must_cleanup_with_gc_fence() {
        // Tests the test util
        let mut engine = TestEngineBuilder::new().build().unwrap();
        must_prewrite_put(&mut engine, b"k", b"v", b"k", 10);
        must_commit(&mut engine, b"k", 10, 20);
        must_cleanup_with_gc_fence(&mut engine, b"k", 20, 0, 30, true);
        let w = must_written(&mut engine, b"k", 10, 20, WriteType::Put);
        assert!(w.has_overlapped_rollback);
        assert_eq!(w.gc_fence.unwrap(), 30.into());
    }

    #[test]
    fn test_cleanup() {
        // Cleanup's logic is mostly similar to rollback, except the TTL check. Tests
        // that not related to TTL check should be covered by other test cases.
        let mut engine = TestEngineBuilder::new().build().unwrap();

        // Shorthand for composing ts.
        let ts = TimeStamp::compose;

        let (k, v) = (b"k", b"v");

        must_prewrite_put(&mut engine, k, v, k, ts(10, 0));
        must_locked(&mut engine, k, ts(10, 0));
        txn_heart_beat::tests::must_success(&mut engine, k, ts(10, 0), 100, 100);
        // Check the last txn_heart_beat has set the lock's TTL to 100.
        txn_heart_beat::tests::must_success(&mut engine, k, ts(10, 0), 90, 100);

        // TTL not expired. Do nothing but returns an error.
        must_err(&mut engine, k, ts(10, 0), ts(20, 0));
        must_locked(&mut engine, k, ts(10, 0));

        // Try to cleanup another transaction's lock. Does nothing.
        must_succeed(&mut engine, k, ts(10, 1), ts(120, 0));
        // If there is no existing lock when cleanup, it may be a pessimistic
        // transaction, so the rollback should be protected.
        must_get_rollback_protected(&mut engine, k, ts(10, 1), true);
        must_locked(&mut engine, k, ts(10, 0));

        // TTL expired. The lock should be removed.
        must_succeed(&mut engine, k, ts(10, 0), ts(120, 0));
        must_unlocked(&mut engine, k);
        // Rollbacks of optimistic transactions need to be protected
        // See: https://github.com/tikv/tikv/issues/16620
        must_get_rollback_protected(&mut engine, k, ts(10, 0), true);
        must_get_rollback_ts(&mut engine, k, ts(10, 0));

        // Rollbacks of primary keys in pessimistic transactions should be protected
        must_acquire_pessimistic_lock(&mut engine, k, k, ts(11, 1), ts(12, 1));
        must_succeed(&mut engine, k, ts(11, 1), ts(120, 0));
        must_get_rollback_protected(&mut engine, k, ts(11, 1), true);

        must_acquire_pessimistic_lock(&mut engine, k, k, ts(13, 1), ts(14, 1));
        must_pessimistic_prewrite_put(
            &mut engine,
            k,
            v,
            k,
            ts(13, 1),
            ts(14, 1),
            DoPessimisticCheck,
        );
        must_succeed(&mut engine, k, ts(13, 1), ts(120, 0));
        must_get_rollback_protected(&mut engine, k, ts(13, 1), true);
    }

    #[test]
    fn test_cleanup_shared_lock_removes_only_matching_holder() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-cleanup";
        let primary = b"primary";
        write_shared_prewrite_locks(&engine, key, primary, &[10, 20]);

        must_succeed(&mut engine, key, 10, 0);
        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_rollback_ts(&mut engine, key, 10);

        must_succeed(&mut engine, key, 20, 0);
        must_unlocked_from_cloud(&mut engine, key);
        must_get_rollback_ts(&mut engine, key, 20);
    }

    #[test]
    fn test_cleanup_rejects_shared_lock_primary() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-cleanup-primary";
        write_shared_prewrite_locks(&engine, key, key, &[10]);

        let err = must_err(&mut engine, key, 10, 0);
        match err {
            MvccError(box ErrorInner::Other(err)) => {
                assert!(
                    err.to_string()
                        .contains("shared lock must not be used as transaction primary")
                );
            }
            err => panic!("unexpected error: {:?}", err),
        }
    }

    #[test]
    fn test_cleanup_rejects_async_commit_shared_lock() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-cleanup-async";
        let primary = b"pk-shared-cleanup-async";
        write_shared_prewrite_locks_with_flags(&engine, key, primary, &[10], true);

        let err = must_err(&mut engine, key, 10, 0);
        match err {
            MvccError(box ErrorInner::Other(err)) => {
                assert!(
                    err.to_string()
                        .contains("shared lock must not use async-commit"),
                    "unexpected error: {:?}",
                    err
                );
            }
            err => panic!("unexpected error: {:?}", err),
        }
    }
}