// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
use tikv_util::{Either, box_err};
use txn_types::{Key, TimeStamp, Write, WriteType};

use crate::storage::{
    Snapshot,
    mvcc::{
        ErrorInner, LockType, MvccTxn, ReleasedLock, Result as MvccResult, SnapshotReader,
        metrics::{
            MVCC_COMMIT_REJECT_BY_BACKUP_TS_COUNTER_VEC, MVCC_CONFLICT_COUNTER,
            MVCC_DUPLICATE_CMD_COUNTER_VEC,
        },
    },
    txn::{
        actions::reject_txn_file_lock,
        commands::{find_mvcc_infos_by_key, find_mvcc_infos_by_key_async},
    },
};

#[maybe_async::both]
pub async fn commit<S: Snapshot>(
    txn: &mut MvccTxn,
    reader: &mut SnapshotReader<S>,
    key: Key,
    commit_ts: TimeStamp,
) -> MvccResult<Option<ReleasedLock>> {
    fail_point!("commit", |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)?,
    };

    let (lock, shared_locks) = match lock_state {
        Some(Either::Left(lock)) if lock.ts == reader.start_ts => (lock, None),
        Some(Either::Right(mut shared_locks)) => {
            txn.mark_shared_lock_lifecycle();
            match shared_locks
                .remove_lock(&reader.start_ts)
                .map_err(crate::storage::mvcc::Error::from)?
            {
                Some(lock) => (lock, Some(shared_locks)),
                None => {
                    return match reader.get_txn_commit_record(&key).await?.info() {
                        Some((_, WriteType::Rollback)) | None => {
                            MVCC_CONFLICT_COUNTER.commit_lock_not_found.inc();
                            info!(
                                "txn conflict (lock not found)";
                                "key" => %key,
                                "start_ts" => reader.start_ts,
                                "commit_ts" => commit_ts,
                            );
                            Err(ErrorInner::TxnLockNotFound {
                                start_ts: reader.start_ts,
                                commit_ts,
                                key: key.into_raw()?,
                            }
                            .into())
                        }
                        Some((_, WriteType::Put))
                        | Some((_, WriteType::Delete))
                        | Some((_, WriteType::Lock)) => {
                            MVCC_DUPLICATE_CMD_COUNTER_VEC.commit.inc();
                            Ok(None)
                        }
                    };
                }
            }
        }
        _ => {
            return match reader.get_txn_commit_record(&key).await?.info() {
                Some((_, WriteType::Rollback)) | None => {
                    MVCC_CONFLICT_COUNTER.commit_lock_not_found.inc();
                    // None: related Rollback has been collapsed.
                    // Rollback: rollback by concurrent transaction.
                    info!(
                        "txn conflict (lock not found)";
                        "key" => %key,
                        "start_ts" => reader.start_ts,
                        "commit_ts" => commit_ts,
                    );
                    Err(ErrorInner::TxnLockNotFound {
                        start_ts: reader.start_ts,
                        commit_ts,
                        key: key.into_raw()?,
                    }
                    .into())
                }
                // Committed by concurrent transaction.
                Some((_, WriteType::Put))
                | Some((_, WriteType::Delete))
                | Some((_, WriteType::Lock)) => {
                    MVCC_DUPLICATE_CMD_COUNTER_VEC.commit.inc();
                    Ok(None)
                }
            };
        }
    };

    reject_txn_file_lock(&lock, &key)?;
    if shared_locks.is_some() && key.is_encoded_from(&lock.primary) {
        return Err(ErrorInner::Other(box_err!(
            "shared lock must not be used as transaction primary"
        ))
        .into());
    }
    if shared_locks.is_some() && lock.use_async_commit {
        return Err(ErrorInner::Other(box_err!("shared lock must not use async-commit")).into());
    }

    let mut min_commit_ts = lock.min_commit_ts;
    let mut reject_by_backup_ts = false;

    // Should not reject async commit.
    if let (Some(false), Some(backup_ts)) = (txn.use_async_commit, txn.backup_ts.as_ref()) {
        if key.is_encoded_from(&lock.primary) {
            let backup_ts = backup_ts.get();
            if min_commit_ts < backup_ts {
                min_commit_ts = backup_ts;
                reject_by_backup_ts = true;
            }
            txn.backup_ts_checked = true;
        }
    }

    // A lock with larger min_commit_ts than current commit_ts can't be committed.
    if commit_ts < min_commit_ts {
        info!(
            "trying to commit with smaller commit_ts than min_commit_ts";
            "key" => %key,
            "start_ts" => reader.start_ts,
            "commit_ts" => commit_ts,
            "min_commit_ts" => min_commit_ts,
            "reject_by_backup_ts" => reject_by_backup_ts,
        );
        if reject_by_backup_ts {
            MVCC_COMMIT_REJECT_BY_BACKUP_TS_COUNTER_VEC.normal.inc();
        }
        return Err(ErrorInner::CommitTsExpired {
            start_ts: reader.start_ts,
            commit_ts,
            key: key.into_raw()?,
            min_commit_ts,
        }
        .into());
    }

    if lock.lock_type == LockType::Pessimistic {
        let writes = match find_mvcc_infos_by_key(reader, &key, TimeStamp::max()).await {
            Ok((_, writes, _)) => Some(writes),
            Err(err) => {
                warn!(
                    "failed to inspect mvcc infos when trying to commit pessimistic lock";
                    "key" => %key,
                    "start_ts" => reader.start_ts,
                    "commit_ts" => commit_ts,
                    "lock" => ?lock,
                    "error" => ?err,
                );
                None
            }
        };
        warn!(
            "rollback a pessimistic lock when trying to commit";
            "key" => %key,
            "start_ts" => reader.start_ts,
            "commit_ts" => commit_ts,
            "lock" => ?lock,
            "writes" => ?writes,
        );
        if let Some(shared_locks) = shared_locks {
            if shared_locks.is_empty() {
                return Ok(txn.delete_shared_lock_wrapper(key, TimeStamp::zero()));
            }
            txn.put_shared_locks(
                key,
                shared_locks,
                crate::storage::mvcc::SharedLockMutation::HolderRemoved,
            );
            return Ok(None);
        }
        return Ok(txn.unlock_key(key, lock.is_pessimistic_txn(), TimeStamp::zero()));
    }

    let mut write = Write::new(
        WriteType::from_lock_type(lock.lock_type).unwrap(),
        reader.start_ts,
        None,
    )
    .set_last_change(lock.last_change_ts, lock.versions_to_last_change)
    .set_txn_source(lock.txn_source);

    for ts in &lock.rollback_ts {
        if *ts == commit_ts {
            write = write.set_overlapped_rollback(true, None);
            break;
        }
    }

    txn.put_write(key.clone(), commit_ts, write.as_ref().to_bytes());
    if let Some(shared_locks) = shared_locks {
        if shared_locks.is_empty() {
            Ok(txn.delete_shared_lock_wrapper(key, commit_ts))
        } else {
            txn.put_shared_locks(
                key,
                shared_locks,
                crate::storage::mvcc::SharedLockMutation::HolderRemoved,
            );
            Ok(None)
        }
    } else {
        Ok(txn.unlock_key(key, lock.is_pessimistic_txn(), commit_ts))
    }
}

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

    use super::*;
    #[cfg(test)]
    use crate::storage::txn::tests::{
        must_acquire_pessimistic_lock, must_acquire_pessimistic_lock_for_large_txn,
        must_pessimistic_prewrite_put, must_prewrite_lock, must_prewrite_put,
        must_prewrite_put_for_large_txn, must_prewrite_put_impl, must_prewrite_put_with_txn_soucre,
    };
    use crate::storage::{
        Engine,
        mvcc::{MvccTxn, tests::*},
    };
    #[cfg(test)]
    use crate::storage::{TestEngineBuilder, TxnStatus, txn::commands::check_txn_status};

    pub fn must_succeed<E: Engine>(
        engine: &mut E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        commit_ts: impl Into<TimeStamp>,
    ) -> Option<ReleasedLock> {
        must_succeed_impl(engine, key, start_ts, commit_ts, None)
    }

    pub fn must_succeed_on_region<E: Engine>(
        engine: &mut E,
        region_id: u64,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        commit_ts: impl Into<TimeStamp>,
    ) -> Option<ReleasedLock> {
        must_succeed_impl(engine, key, start_ts, commit_ts, Some(region_id))
    }

    fn must_succeed_impl<E: Engine>(
        engine: &mut E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        commit_ts: impl Into<TimeStamp>,
        region_id: Option<u64>,
    ) -> Option<ReleasedLock> {
        let mut ctx = Context::default();
        if let Some(region_id) = region_id {
            ctx.region_id = region_id;
        }
        let snap_ctx = SnapContext {
            pb_ctx: &ctx,
            ..Default::default()
        };
        let snapshot = engine.snapshot(snap_ctx).unwrap();
        let start_ts = start_ts.into();
        let cm = ConcurrencyManager::new(start_ts);
        let mut txn = MvccTxn::new(start_ts, cm);
        let mut reader = SnapshotReader::new(start_ts, snapshot, true);
        let released_lock =
            commit(&mut txn, &mut reader, Key::from_raw(key), commit_ts.into()).unwrap();
        let shared_lock_lifecycle = txn.has_shared_lock_lifecycle();
        let mut data = WriteData::from_modifies(txn.into_modifies());
        data.extra.req_type = ReqType::Commit;
        data.set_shared_lock_lifecycle(shared_lock_lifecycle);
        write(engine, &ctx, data);
        released_lock
    }

    pub fn must_err<E: Engine>(
        engine: &mut E,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        commit_ts: impl Into<TimeStamp>,
    ) -> crate::storage::mvcc::Error {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let start_ts = start_ts.into();
        let cm = ConcurrencyManager::new(start_ts);
        let mut txn = MvccTxn::new(start_ts, cm);
        let mut reader = SnapshotReader::new(start_ts, snapshot, true);
        commit(&mut txn, &mut reader, Key::from_raw(key), commit_ts.into()).unwrap_err()
    }

    #[cfg(test)]
    fn write_shared_prewrite_locks<E: Engine>(
        engine: &E,
        key: &[u8],
        primary: &[u8],
        holders: &[(u64, 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, u64)],
        use_async_commit: bool,
    ) {
        let mut shared_locks = txn_types::SharedLocks::new();
        for (start_ts, min_commit_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,
                (*min_commit_ts).into(),
            );
            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 = WriteData::from_modifies(vec![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());
    }

    #[cfg(test)]
    fn must_commit_shared_lock(
        engine: &mut TestKvEngine,
        key: &[u8],
        start_ts: impl Into<TimeStamp>,
        commit_ts: impl Into<TimeStamp>,
    ) {
        let start_ts = start_ts.into();
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let mut reader = SnapshotReader::new(start_ts, snapshot, true);
        let (actual_commit_ts, write_type) = reader
            .get_txn_commit_record(&Key::from_raw(key))
            .unwrap()
            .info()
            .unwrap();
        assert_eq!(actual_commit_ts, commit_ts.into());
        assert_eq!(write_type, WriteType::Lock);
    }

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

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

        // Shortcuts
        let ts = TimeStamp::compose;
        let uncommitted = |ttl, min_commit_ts| {
            move |s| {
                if let TxnStatus::Uncommitted { lock, .. } = s {
                    lock.ttl == ttl && lock.min_commit_ts == min_commit_ts
                } else {
                    false
                }
            }
        };

        must_prewrite_put_for_large_txn(&mut engine, k, v, k, ts(10, 0), 100, 0);
        check_txn_status::tests::must_success(
            &mut engine,
            k,
            ts(10, 0),
            ts(20, 0),
            ts(20, 0),
            true,
            false,
            false,
            uncommitted(100, ts(20, 1)),
        );
        // The min_commit_ts should be ts(20, 1)
        must_err(&mut engine, k, ts(10, 0), ts(15, 0));
        must_err(&mut engine, k, ts(10, 0), ts(20, 0));
        must_succeed(&mut engine, k, ts(10, 0), ts(20, 1));

        must_prewrite_put_for_large_txn(&mut engine, k, v, k, ts(30, 0), 100, 0);
        check_txn_status::tests::must_success(
            &mut engine,
            k,
            ts(30, 0),
            ts(40, 0),
            ts(40, 0),
            true,
            false,
            false,
            uncommitted(100, ts(40, 1)),
        );
        must_succeed(&mut engine, k, ts(30, 0), ts(50, 0));

        // If the min_commit_ts of the pessimistic lock is greater than prewrite's, use
        // it.
        must_acquire_pessimistic_lock_for_large_txn(&mut engine, k, k, ts(60, 0), ts(60, 0), 100);
        check_txn_status::tests::must_success(
            &mut engine,
            k,
            ts(60, 0),
            ts(70, 0),
            ts(70, 0),
            true,
            false,
            false,
            uncommitted(100, ts(70, 1)),
        );
        must_prewrite_put_impl(
            &mut engine,
            k,
            v,
            k,
            &None,
            ts(60, 0),
            DoPessimisticCheck,
            50,
            ts(60, 0),
            1,
            ts(60, 1),
            TimeStamp::zero(),
            false,
            kvproto::kvrpcpb::Assertion::None,
            kvproto::kvrpcpb::AssertionLevel::Off,
        );
        // The min_commit_ts is ts(70, 0) other than ts(60, 1) in prewrite request.
        must_large_txn_locked(&mut engine, k, ts(60, 0), 100, ts(70, 1), false);
        must_err(&mut engine, k, ts(60, 0), ts(65, 0));
        must_succeed(&mut engine, k, ts(60, 0), ts(80, 0));
    }

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

        let key = b"k-shared-min-commit-ts";
        let primary = b"pk-shared-min-commit-ts";
        write_shared_prewrite_locks(&engine, key, primary, &[(10, 20)]);

        must_err(&mut engine, key, 10, 19);
        let shared_locks = must_shared_locks(&mut engine, key);
        assert!(shared_locks.contains_start_ts(10.into()));

        must_succeed(&mut engine, key, 10, 20);
        must_unlocked_from_cloud(&mut engine, key);
        must_commit_shared_lock(&mut engine, key, 10, 20);
    }

    #[test]
    fn test_commit_shared_lock_emits_upgrade_hint_for_remaining_holder() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-commit-upgrade-hint";
        write_shared_prewrite_locks(&engine, key, b"primary", &[(10, 20), (20, 30)]);
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let mut reader = SnapshotReader::new(20.into(), snapshot, true);
        let mut txn = MvccTxn::new(20.into(), ConcurrencyManager::new(TimeStamp::zero()));

        commit(&mut txn, &mut reader, Key::from_raw(key), 30.into()).unwrap();

        let hints = txn.take_shared_lock_upgrade_hints();
        assert_eq!(hints.len(), 1);
        assert_eq!(hints[0].key, Key::from_raw(key));
        assert_eq!(hints[0].sole_holder_start_ts, 10.into());
    }

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

        let err = must_err(&mut engine, key, 10, 20);
        match err {
            crate::storage::mvcc::Error(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_commit_rejects_async_commit_shared_lock() {
        let mut engine = TestKvEngine::new().unwrap();
        let key = b"k-shared-commit-async";
        let primary = b"pk-shared-commit-async";
        write_shared_prewrite_locks_with_flags(&engine, key, primary, &[(10, 20)], true);

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

    #[test]
    #[ignore = "CSE incompatibility; see #6158"]
    fn test_inherit_last_change_info_from_lock() {
        let mut engine = TestEngineBuilder::new().build().unwrap();

        let k = b"k";
        must_prewrite_put(&mut engine, k, b"v1", k, 5);
        must_succeed(&mut engine, k, 5, 10);

        // WriteType is Lock
        must_prewrite_lock(&mut engine, k, k, 15);
        let lock = must_locked(&mut engine, k, 15);
        assert_eq!(lock.last_change_ts, 10.into());
        assert_eq!(lock.versions_to_last_change, 1);
        must_succeed(&mut engine, k, 15, 20);
        let write = must_written(&mut engine, k, 15, 20, WriteType::Lock);
        assert_eq!(write.last_change_ts, 10.into());
        assert_eq!(write.versions_to_last_change, 1);

        // WriteType is Put
        must_prewrite_put(&mut engine, k, b"v2", k, 25);
        let lock = must_locked(&mut engine, k, 25);
        assert!(lock.last_change_ts.is_zero());
        assert_eq!(lock.versions_to_last_change, 0);
        must_succeed(&mut engine, k, 25, 30);
        let write = must_written(&mut engine, k, 25, 30, WriteType::Put);
        assert!(write.last_change_ts.is_zero());
        assert_eq!(write.versions_to_last_change, 0);
    }

    #[test]
    #[ignore = "CSE incompatibility; see #6158"]
    fn test_2pc_with_txn_source() {
        for source in [0x1, 0x85] {
            let mut engine = TestEngineBuilder::new().build().unwrap();

            let k = b"k";
            // WriteType is Put
            must_prewrite_put_with_txn_soucre(&mut engine, k, b"v2", k, 25, source);
            let lock = must_locked(&mut engine, k, 25);
            assert_eq!(lock.txn_source, source);
            must_succeed(&mut engine, k, 25, 30);
            let write = must_written(&mut engine, k, 25, 30, WriteType::Put);
            assert_eq!(write.txn_source, source);
        }
    }

    #[test]
    #[ignore = "CSE incompatibility; see #6158"]
    fn test_commit_rollback_pessimistic_lock() {
        let mut engine = TestEngineBuilder::new().build().unwrap();

        let k1 = b"k1";
        let k2 = b"k2";

        must_acquire_pessimistic_lock(&mut engine, k1, k1, 10, 10);
        must_acquire_pessimistic_lock(&mut engine, k2, k1, 10, 10);
        must_pessimistic_prewrite_put(&mut engine, k1, b"v1", k1, 10, 10, DoPessimisticCheck);

        let released_lock = must_succeed(&mut engine, k1, 10, 20).unwrap();
        assert_eq!(released_lock.key, Key::from_raw(k1));
        assert_eq!(released_lock.start_ts, 10.into());
        assert_eq!(released_lock.commit_ts, 20.into());

        let released_lock = must_succeed(&mut engine, k2, 10, 20).unwrap();
        assert_eq!(released_lock.key, Key::from_raw(k2));
        assert_eq!(released_lock.start_ts, 10.into());
        assert_eq!(released_lock.commit_ts, TimeStamp::zero());

        must_written(&mut engine, k1, 10, 20, WriteType::Put);
        must_not_have_write(&mut engine, k2, 20);
        must_not_have_write(&mut engine, k2, 10);
    }
}