// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
// #[PerformanceCriticalPath]
use txn_types::{Key, ReqType, TxnExtra};

use crate::storage::{
    ProcessResult, Snapshot, TxnStatus,
    kv::WriteData,
    lock_manager::LockManager,
    mvcc::{MvccTxn, SnapshotReader},
    txn::{
        Error, ErrorInner, Result,
        commands::{
            Command, CommandExt, ReaderWithStats, ReleasedLocks, ResponsePolicy, TypedCommand,
            WriteCommand, WriteContext, WriteResult,
        },
        commit, commit_async,
    },
};

command! {
    /// Commit the transaction that started at `lock_ts`.
    ///
    /// This should be following a [`Prewrite`](Command::Prewrite).
    Commit:
        cmd_ty => TxnStatus,
        display => "kv::command::commit {:?} {} -> {} | {:?}", (keys, lock_ts, commit_ts, ctx),
        content => {
            /// The keys affected.
            keys: Vec<Key>,
            /// The lock timestamp.
            lock_ts: txn_types::TimeStamp,
            /// The commit timestamp.
            commit_ts: txn_types::TimeStamp,
            /// The commit is using async commit.
            use_async_commit: bool,
            /// Used in file based transaction.
            is_txn_file: bool,
        }
}

impl CommandExt for Commit {
    ctx!();
    tag!(commit);
    request_type!(KvCommit);
    ts!(commit_ts);
    write_bytes!(keys: multiple);
    gen_lock!(keys: multiple);
    can_build_txn_file!();
}

#[maybe_async::async_trait]
impl<S: Snapshot + 'static, L: LockManager> WriteCommand<S, L> for Commit {
    #[maybe_async]
    async fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
        if self.commit_ts <= self.lock_ts {
            return Err(Error::from(ErrorInner::InvalidTxnTso {
                start_ts: self.lock_ts,
                commit_ts: self.commit_ts,
            }));
        }
        let mut txn = MvccTxn::new_with_backup_ts(
            self.lock_ts,
            context.concurrency_manager,
            self.use_async_commit,
        );
        let mut reader = ReaderWithStats::new(
            SnapshotReader::new_with_ctx(self.lock_ts, snapshot, &self.ctx),
            context.statistics,
        );

        fail::fail_point!("txn::commit_process_write_lock_not_found", |_| {
            // To generate a retryable error for testing purpose.
            Err(self.txn_lock_not_found_error())
        });

        let rows = self.keys.len();
        // Pessimistic txn needs key_hashes to wake up waiters
        let mut released_locks = ReleasedLocks::new();
        let mut has_stale_pessimistic_lock_cleanup = false;
        for k in self.keys {
            let released_lock = commit(&mut txn, &mut reader, k, self.commit_ts).await?;
            if matches!(released_lock.as_ref(), Some(lock) if lock.commit_ts.is_zero()) {
                has_stale_pessimistic_lock_cleanup = true;
            }
            released_locks.push(released_lock);
        }

        let pr = ProcessResult::TxnStatus {
            txn_status: TxnStatus::committed(self.commit_ts),
        };
        let shared_lock_upgrade_hints = txn.take_shared_lock_upgrade_hints();
        let backup_ts_checked = txn.take_checked_backup_ts();
        let shared_lock_lifecycle = txn.has_shared_lock_lifecycle();
        let mut write_data =
            WriteData::new(txn.into_modifies(), TxnExtra::default(), backup_ts_checked);
        // `req_type` is used as the custom raft-log encoding selector here,
        // not as the original RPC identity. A kv_commit batch with stale
        // pessimistic-lock cleanup is routed through ResolveLock so it can
        // reuse the existing delete-lock-only rollback format without adding a
        // new persisted format. See upstream tikv/tikv#14557.
        write_data.set_req_type(if has_stale_pessimistic_lock_cleanup {
            ReqType::ResolveLock
        } else {
            ReqType::Commit
        });
        write_data.set_shared_lock_lifecycle(shared_lock_lifecycle);
        write_data.set_allowed_on_disk_almost_full();

        fail::fail_point!("txn::commit_process_write_finish");

        Ok(WriteResult {
            ctx: self.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![(self.lock_ts, self.commit_ts)],
        })
    }
}

#[cfg(feature = "failpoints")]
impl Commit {
    fn txn_lock_not_found_error(&self) -> Error {
        let key = self
            .keys
            .first()
            .and_then(|k| k.to_raw().ok())
            .unwrap_or_default();
        Error::from_mvcc(crate::storage::mvcc::ErrorInner::TxnLockNotFound {
            start_ts: self.lock_ts,
            commit_ts: self.commit_ts,
            key,
        })
    }
}

#[cfg(test)]
mod tests {
    use concurrency_manager::ConcurrencyManager;
    use engine_traits::{CF_LOCK, CF_WRITE};
    use kvproto::kvrpcpb::{
        Context, ExtraOp, PrewriteRequestPessimisticAction::DoPessimisticCheck,
    };

    use super::*;
    use crate::storage::{
        Engine, MockLockManager, Statistics, TestEngineBuilder,
        kv::Modify,
        txn::{
            commands::{WriteContext, test_util as command_test_util},
            tests::{must_acquire_pessimistic_lock, must_pessimistic_prewrite_put},
            txn_status_cache::TxnStatusCache,
        },
    };

    fn process_commit_write<E: Engine>(
        engine: &mut E,
        keys: Vec<Key>,
        lock_ts: u64,
        commit_ts: u64,
    ) -> WriteResult {
        let snapshot = engine.snapshot(Default::default()).unwrap();
        let cmd = Commit::new(
            keys,
            lock_ts.into(),
            commit_ts.into(),
            false,
            false,
            Context::default(),
        );
        let lock_mgr = MockLockManager::new();
        let txn_status_cache = TxnStatusCache::new_for_test();
        let mut statistics = Statistics::default();

        cmd.cmd
            .process_write(
                snapshot,
                WriteContext {
                    lock_mgr: &lock_mgr,
                    concurrency_manager: ConcurrencyManager::new(lock_ts.into()),
                    extra_op: ExtraOp::Noop,
                    extra_cf_scan_mode: Default::default(),
                    statistics: &mut statistics,
                    async_apply_prewrite: false,
                    txn_status_cache: &txn_status_cache,
                },
            )
            .unwrap()
    }

    #[test]
    fn test_process_write_uses_commit_req_type_for_duplicate_commit() {
        let mut engine = TestEngineBuilder::new().build().unwrap();
        let key = b"k1";

        must_acquire_pessimistic_lock(&mut engine, key, key, 10, 10);
        must_pessimistic_prewrite_put(&mut engine, key, b"v1", key, 10, 10, DoPessimisticCheck);
        command_test_util::commit(
            &mut engine,
            &mut Statistics::default(),
            vec![Key::from_raw(key)],
            10,
            20,
        )
        .unwrap();

        let result = process_commit_write(&mut engine, vec![Key::from_raw(key)], 10, 20);
        assert!(result.to_be_write.modifies.is_empty());
        assert_eq!(result.to_be_write.extra.req_type, ReqType::Commit);
    }

    #[test]
    fn test_process_write_uses_commit_req_type_for_normal_commit() {
        let mut engine = TestEngineBuilder::new().build().unwrap();
        let key = b"k1";

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

        let result = process_commit_write(&mut engine, vec![Key::from_raw(key)], 10, 20);
        assert_eq!(result.to_be_write.modifies.len(), 2);
        assert_eq!(result.to_be_write.extra.req_type, ReqType::Commit);
    }

    #[test]
    fn test_process_write_uses_resolve_lock_req_type_for_delete_only_commit() {
        let mut engine = TestEngineBuilder::new().build().unwrap();
        let key = b"k1";

        must_acquire_pessimistic_lock(&mut engine, key, key, 10, 10);

        let result = process_commit_write(&mut engine, vec![Key::from_raw(key)], 10, 20);
        assert_eq!(result.to_be_write.modifies.len(), 1);
        assert_eq!(result.to_be_write.extra.req_type, ReqType::ResolveLock);
        assert!(matches!(
            result.to_be_write.modifies.first(),
            Some(Modify::Delete(cf, modify_key))
                if *cf == CF_LOCK && modify_key == &Key::from_raw(key)
        ));
    }

    #[test]
    fn test_process_write_uses_resolve_lock_req_type_for_mixed_commit() {
        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 result = process_commit_write(
            &mut engine,
            vec![Key::from_raw(k1), Key::from_raw(k2)],
            10,
            20,
        );

        assert_eq!(result.to_be_write.extra.req_type, ReqType::ResolveLock);
        assert!(matches!(
            result.to_be_write.modifies.first(),
            Some(Modify::Put(cf, key, _))
                if *cf == CF_WRITE && key == &Key::from_raw(k1).append_ts(20.into())
        ));
        assert!(matches!(
            result.to_be_write.modifies.get(1),
            Some(Modify::Delete(cf, key)) if *cf == CF_LOCK && key == &Key::from_raw(k1)
        ));
        assert!(matches!(
            result.to_be_write.modifies.get(2),
            Some(Modify::Delete(cf, key)) if *cf == CF_LOCK && key == &Key::from_raw(k2)
        ));
    }
}