// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
use kvproto::kvrpcpb::{ExtraOp, WriteConflictReason};
use resource_metering::record_network_out_bytes;
use tikv_kv::Modify;
use txn_types::{Key, OldValues, ReqType, TimeStamp, TxnExtra, insert_old_value_if_resolved};

use crate::storage::{
    Error as StorageError, PessimisticLockKeyResult, ProcessResult, Result as StorageResult,
    Snapshot,
    kv::WriteData,
    lock_manager::{LockManager, WaitTimeout},
    mvcc::{
        Error as MvccError, ErrorInner as MvccErrorInner, MvccTxn, SharedLocksShrinkReason,
        SnapshotReader,
    },
    txn::{
        Error, ErrorInner, Result, acquire_pessimistic_lock, acquire_pessimistic_lock_async,
        commands::{
            Command, CommandExt, ReaderWithStats, ReleasedLocks, ResponsePolicy, TypedCommand,
            WriteCommand, WriteContext, WriteResult, WriteResultLockInfo,
        },
    },
    types::{PessimisticLockParameters, PessimisticLockResults},
};

command! {
    /// Acquire a Pessimistic lock on the keys.
    ///
    /// This can be rolled back with a [`PessimisticRollback`](Command::PessimisticRollback) command.
    AcquirePessimisticLock:
        cmd_ty => StorageResult<PessimisticLockResults>,
        display => "kv::command::acquirepessimisticlock keys({:?}) @ {} {} {} {:?} {} {} {} | {:?}",
        (keys, start_ts, lock_ttl, for_update_ts, wait_timeout, min_commit_ts, check_existence, lock_only_if_exists, ctx),
        content => {
            /// The set of keys to lock.
            /// `(Key, should_not_exist, is_shared_lock)`
            keys: Vec<(Key, bool, bool)>,
            /// The primary lock. Secondary locks (from `keys`) will refer to the primary lock.
            primary: Vec<u8>,
            /// The transaction timestamp.
            start_ts: TimeStamp,
            /// The Time To Live of the lock, in milliseconds
            lock_ttl: u64,
            is_first_lock: bool,
            for_update_ts: TimeStamp,
            /// Time to wait for lock released in milliseconds when encountering locks.
            wait_timeout: Option<WaitTimeout>,
            /// If it is true, TiKV will return values of the keys if no error, so TiDB can cache the values for
            /// later read in the same transaction.
            return_values: bool,
            min_commit_ts: TimeStamp,
            check_existence: bool,
            lock_only_if_exists: bool,
            allow_lock_with_conflict: bool,
        }
}

impl CommandExt for AcquirePessimisticLock {
    ctx!();
    tag!(acquire_pessimistic_lock);
    request_type!(KvPessimisticLock);
    ts!(start_ts);
    property!(can_be_pipelined);

    fn write_bytes(&self) -> usize {
        self.keys
            .iter()
            .map(|(key, ..)| key.as_encoded().len())
            .sum()
    }

    gen_lock!(keys: multiple(|x| &x.0));
}

#[maybe_async::async_trait]
impl<S: Snapshot + 'static, L: LockManager> WriteCommand<S, L> for AcquirePessimisticLock {
    #[maybe_async]
    async fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
        if self.allow_lock_with_conflict {
            // Fair locking (allow_lock_with_conflict) is not supported on cloud-engine.
            return Err(Error::from(ErrorInner::Other(box_err!(
                "allow_lock_with_conflict is not supported on cloud-engine"
            ))));
        }

        let (start_ts, ctx, keys) = (self.start_ts, self.ctx, self.keys);
        let mut txn = MvccTxn::new(start_ts, context.concurrency_manager);
        let mut reader = ReaderWithStats::new(
            SnapshotReader::new_with_ctx_and_scan_mode(
                start_ts,
                snapshot,
                &ctx,
                context.extra_cf_scan_mode,
            ),
            context.statistics,
        );

        let total_keys = keys.len();
        let mut res = PessimisticLockResults::with_capacity(total_keys);
        let mut encountered_locks = vec![];
        let mut shrink_error = None;
        let need_old_value = context.extra_op == ExtraOp::ReadOldValue;
        let mut old_values = OldValues::default();
        for (k, should_not_exist, is_shared_lock) in keys {
            if is_shared_lock && k.is_encoded_from(&self.primary) {
                return Err(Error::from(ErrorInner::Other(box_err!(
                    "shared pessimistic lock must not be used as transaction primary"
                ))));
            }
            match acquire_pessimistic_lock(
                &mut txn,
                &mut reader,
                k.clone(),
                &self.primary,
                should_not_exist,
                self.lock_ttl,
                self.for_update_ts,
                self.return_values,
                self.check_existence,
                self.min_commit_ts,
                need_old_value,
                self.lock_only_if_exists,
                self.allow_lock_with_conflict,
                is_shared_lock,
            )
            .await
            {
                Ok((key_res, old_value)) => {
                    res.push(key_res);
                    // MutationType is unknown in AcquirePessimisticLock stage.
                    insert_old_value_if_resolved(&mut old_values, k, txn.start_ts, old_value, None);
                }
                Err(MvccError(box MvccErrorInner::KeyIsLocked(lock_info))) => {
                    let request_parameters = PessimisticLockParameters {
                        pb_ctx: ctx.clone(),
                        primary: self.primary.clone(),
                        start_ts,
                        lock_ttl: self.lock_ttl,
                        for_update_ts: self.for_update_ts,
                        wait_timeout: self.wait_timeout,
                        return_values: self.return_values,
                        min_commit_ts: self.min_commit_ts,
                        check_existence: self.check_existence,
                        is_first_lock: self.is_first_lock,
                        lock_only_if_exists: self.lock_only_if_exists,
                        allow_lock_with_conflict: self.allow_lock_with_conflict,
                    };
                    let lock_info = WriteResultLockInfo::new(
                        lock_info,
                        request_parameters,
                        k,
                        should_not_exist,
                        is_shared_lock,
                    );
                    encountered_locks.push(lock_info);
                    // Do not lock previously succeeded keys.
                    txn.clear();
                    res.0.clear();
                    res.push(PessimisticLockKeyResult::Waiting);
                    break;
                }
                Err(MvccError(box MvccErrorInner::SharedLocksNeedShrink {
                    mut shared_locks,
                    reason,
                })) => {
                    let locked_raw_key = k.to_raw()?;
                    shared_locks.set_shrink_only();
                    txn.clear();
                    txn.put_shared_lock_flags(k, &shared_locks);
                    old_values.clear();
                    res.0.clear();
                    shrink_error = Some(match reason {
                        SharedLocksShrinkReason::NonHolderExclusive => StorageError::from(
                            Error::from(MvccError::from(MvccErrorInner::KeyIsLocked(
                                shared_locks.into_lock_info(locked_raw_key),
                            ))),
                        ),
                        SharedLocksShrinkReason::MultiHolderUpgrade => StorageError::from(
                            Error::from(MvccError::from(MvccErrorInner::WriteConflict {
                                start_ts,
                                conflict_start_ts: TimeStamp::zero(),
                                conflict_commit_ts: TimeStamp::zero(),
                                key: locked_raw_key,
                                primary: self.primary.clone(),
                                reason: WriteConflictReason::PessimisticRetry,
                            })),
                        ),
                    });
                    break;
                }
                Err(e) => return Err(Error::from(e)),
            }
        }

        let new_acquired_locks = txn.take_new_locks();
        let shared_lock_upgrade_hints = txn.take_shared_lock_upgrade_hints();
        let shared_lock_lifecycle = txn.has_shared_lock_lifecycle();
        let modifies = txn.into_modifies();

        let mut res = Ok(res);

        // If encountered lock and `wait_timeout` is `None` (which means no wait),
        // return error directly here.
        if !encountered_locks.is_empty() && self.wait_timeout.is_none() {
            // Mind the difference of the protocols of legacy requests and resumable
            // requests. For resumable requests (allow_lock_with_conflict ==
            // true), key errors are considered key by key instead of for the
            // whole request.
            let lock_info = encountered_locks.drain(..).next().unwrap().lock_info_pb;
            let err = StorageError::from(Error::from(MvccError::from(
                MvccErrorInner::KeyIsLocked(lock_info),
            )));
            if self.allow_lock_with_conflict {
                res.as_mut().unwrap().0[0] = PessimisticLockKeyResult::Failed(err.into())
            } else {
                res = Err(err)
            }
        }
        if let Some(err) = shrink_error {
            res = Err(err);
        }

        let rows = if res.is_ok() { total_keys } else { 0 };

        record_network_out_bytes(res.as_ref().map_or(0, |v| v.estimate_resp_size()));
        let pr = ProcessResult::PessimisticLockRes { res };

        // Probe hook for old_value testing (always compiled for test package)
        if tikv_kv::PROBE_OLD_VALUES_IN_TEST.load(std::sync::atomic::Ordering::Relaxed) {
            tikv_kv::populate_old_values_probe_cache(&old_values, &[]);
        }

        let mut to_be_write = make_write_data(modifies, old_values, ReqType::PessimisticLock);
        to_be_write.set_shared_lock_lifecycle(shared_lock_lifecycle);

        Ok(WriteResult {
            ctx,
            to_be_write,
            rows,
            pr,
            lock_info: encountered_locks,
            released_locks: ReleasedLocks::new(),
            new_acquired_locks,
            shared_lock_upgrade_hints,
            lock_guards: vec![],
            response_policy: ResponsePolicy::OnProposed,
            known_txn_status: vec![],
        })
    }
}

pub(super) fn make_write_data(
    modifies: Vec<Modify>,
    old_values: OldValues,
    req_type: ReqType,
) -> WriteData {
    if !modifies.is_empty() {
        let extra = TxnExtra {
            old_values,
            // One pc status is unknown in AcquirePessimisticLock stage.
            one_pc: false,
            for_flashback: false,
            req_type,
            shared_lock_lifecycle: false,
        };
        WriteData::new(modifies, extra, None)
    } else {
        WriteData::default()
    }
}

#[cfg(test)]
mod tests {
    use cloud_server::mock_kv_engine::TestKvEngine;
    use concurrency_manager::ConcurrencyManager;
    use kvproto::kvrpcpb::{Context, ExtraOp};
    use txn_types::TimeStamp;

    use super::*;
    use crate::storage::{
        ProcessResult, TestEngineBuilder,
        errors::ErrorInner as StorageErrorInner,
        kv::{Engine, Statistics},
        lock_manager::MockLockManager,
        mvcc::{Error as MvccError, ErrorInner as MvccErrorInner},
        txn::{
            Error as TxnError, ErrorInner as TxnErrorInner,
            commands::WriteContext,
            tests::{must_commit, must_prewrite_put},
            txn_status_cache::TxnStatusCache,
        },
    };

    fn process_acquire_pessimistic_lock_with_options(
        engine: &mut impl Engine,
        key: &[u8],
        start_ts: u64,
        for_update_ts: u64,
        return_values: bool,
        check_existence: bool,
        allow_lock_with_conflict: bool,
        is_shared_lock: bool,
    ) -> Result<WriteResult> {
        process_acquire_pessimistic_lock_with_primary_and_options(
            engine,
            key,
            b"pk",
            start_ts,
            for_update_ts,
            return_values,
            check_existence,
            allow_lock_with_conflict,
            is_shared_lock,
        )
    }

    fn process_acquire_pessimistic_lock_with_primary_and_options(
        engine: &mut impl Engine,
        key: &[u8],
        primary: &[u8],
        start_ts: u64,
        for_update_ts: u64,
        return_values: bool,
        check_existence: bool,
        allow_lock_with_conflict: bool,
        is_shared_lock: bool,
    ) -> Result<WriteResult> {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let lock_mgr = MockLockManager::new();
        let txn_status_cache = TxnStatusCache::new_for_test();
        let mut statistics = Statistics::default();
        AcquirePessimisticLock::new(
            vec![(Key::from_raw(key), false, is_shared_lock)],
            primary.to_vec(),
            start_ts.into(),
            1000,
            false,
            for_update_ts.into(),
            None,
            return_values,
            TimeStamp::zero(),
            check_existence,
            false,
            allow_lock_with_conflict,
            Context::default(),
        )
        .cmd
        .process_write(
            snapshot,
            WriteContext {
                lock_mgr: &lock_mgr,
                concurrency_manager: ConcurrencyManager::new(TimeStamp::zero()),
                extra_op: ExtraOp::Noop,
                extra_cf_scan_mode: Default::default(),
                statistics: &mut statistics,
                async_apply_prewrite: false,
                txn_status_cache: &txn_status_cache,
            },
        )
    }

    fn process_acquire_pessimistic_lock(
        engine: &mut impl Engine,
        key: &[u8],
        start_ts: u64,
        allow_lock_with_conflict: bool,
        is_shared_lock: bool,
    ) -> Result<WriteResult> {
        process_acquire_pessimistic_lock_with_options(
            engine,
            key,
            start_ts,
            start_ts,
            false,
            false,
            allow_lock_with_conflict,
            is_shared_lock,
        )
    }

    fn pessimistic_lock_results(result: &WriteResult) -> &PessimisticLockResults {
        match &result.pr {
            ProcessResult::PessimisticLockRes { res: Ok(res) } => res,
            pr => panic!("unexpected process result: {:?}", pr),
        }
    }

    #[test]
    fn test_shared_lock_acquire_marks_lifecycle() {
        let mut engine = TestKvEngine::new().unwrap();

        let result = process_acquire_pessimistic_lock(&mut engine, b"k", 5, false, true).unwrap();
        assert!(result.to_be_write.extra.shared_lock_lifecycle);

        let result = process_acquire_pessimistic_lock(&mut engine, b"k2", 6, false, false).unwrap();
        assert!(!result.to_be_write.extra.shared_lock_lifecycle);
    }

    #[test]
    fn test_shrink_only_update_marks_shared_lock_lifecycle() {
        let mut engine = TestKvEngine::new().unwrap();
        let result = process_acquire_pessimistic_lock(&mut engine, b"k", 5, false, true).unwrap();
        engine
            .write(&Context::default(), result.to_be_write)
            .unwrap();

        let result = process_acquire_pessimistic_lock(&mut engine, b"k", 6, false, false).unwrap();
        assert!(result.to_be_write.extra.shared_lock_lifecycle);
        assert!(matches!(
            result.pr,
            ProcessResult::PessimisticLockRes {
                res: Err(StorageError(box StorageErrorInner::Txn(TxnError(box TxnErrorInner::Mvcc(
                    MvccError(box MvccErrorInner::KeyIsLocked(_))
                )))))
            }
        ));
    }

    #[test]
    fn test_multi_holder_upgrade_returns_pessimistic_retry() {
        let mut engine = TestKvEngine::new().unwrap();
        for start_ts in [5, 6] {
            let result =
                process_acquire_pessimistic_lock(&mut engine, b"k", start_ts, false, true).unwrap();
            engine
                .write(&Context::default(), result.to_be_write)
                .unwrap();
        }

        let result = process_acquire_pessimistic_lock(&mut engine, b"k", 5, false, false).unwrap();
        assert!(result.to_be_write.extra.shared_lock_lifecycle);
        assert!(matches!(
            result.pr,
            ProcessResult::PessimisticLockRes {
                res: Err(StorageError(box StorageErrorInner::Txn(TxnError(box TxnErrorInner::Mvcc(
                    MvccError(box MvccErrorInner::WriteConflict {
                        reason: WriteConflictReason::PessimisticRetry,
                        ..
                    })
                )))))
            }
        ));
    }

    #[test]
    fn test_shared_lock_retry_preserves_return_value_and_existence() {
        let mut engine = TestKvEngine::new().unwrap();

        must_prewrite_put(&mut engine, b"k-value", b"v1", b"k-value", 5);
        must_commit(&mut engine, b"k-value", 5, 20);

        let result = process_acquire_pessimistic_lock_with_options(
            &mut engine,
            b"k-value",
            30,
            30,
            true,
            false,
            false,
            true,
        )
        .unwrap();
        pessimistic_lock_results(&result).0[0].assert_value(Some(b"v1"));
        engine
            .write(&Context::default(), result.to_be_write)
            .unwrap();

        let retry = process_acquire_pessimistic_lock_with_options(
            &mut engine,
            b"k-value",
            30,
            30,
            true,
            false,
            false,
            true,
        )
        .unwrap();
        pessimistic_lock_results(&retry).0[0].assert_value(Some(b"v1"));

        must_prewrite_put(&mut engine, b"k-exists", b"v2", b"k-exists", 6);
        must_commit(&mut engine, b"k-exists", 6, 21);

        let result = process_acquire_pessimistic_lock_with_options(
            &mut engine,
            b"k-exists",
            40,
            40,
            false,
            true,
            false,
            true,
        )
        .unwrap();
        pessimistic_lock_results(&result).0[0].assert_existence(true);
        engine
            .write(&Context::default(), result.to_be_write)
            .unwrap();

        let retry = process_acquire_pessimistic_lock_with_options(
            &mut engine,
            b"k-exists",
            40,
            40,
            false,
            true,
            false,
            true,
        )
        .unwrap();
        pessimistic_lock_results(&retry).0[0].assert_existence(true);
    }

    #[test]
    fn test_shared_lock_acquire_rejects_allow_lock_with_conflict() {
        let mut engine = TestEngineBuilder::new().build().unwrap();
        let err = match process_acquire_pessimistic_lock(&mut engine, b"k", 5, true, true) {
            Ok(_) => panic!("expected allow_lock_with_conflict to be rejected"),
            Err(err) => err,
        };
        match err {
            Error(box ErrorInner::Other(err)) => {
                assert!(
                    err.to_string()
                        .contains("allow_lock_with_conflict is not supported on cloud-engine"),
                    "unexpected error: {:?}",
                    err
                );
            }
            err => panic!("unexpected error: {:?}", err),
        }
    }

    #[test]
    fn test_shared_lock_acquire_rejects_primary_key() {
        let mut engine = TestKvEngine::new().unwrap();
        let err = match process_acquire_pessimistic_lock_with_primary_and_options(
            &mut engine,
            b"k",
            b"k",
            5,
            5,
            false,
            false,
            false,
            true,
        ) {
            Ok(_) => panic!("expected shared pessimistic lock primary to be rejected"),
            Err(err) => err,
        };
        match err {
            Error(box ErrorInner::Other(err)) => {
                assert!(
                    err.to_string().contains(
                        "shared pessimistic lock must not be used as transaction primary"
                    ),
                    "unexpected error: {:?}",
                    err
                );
            }
            err => panic!("unexpected error: {:?}", err),
        }
    }
}