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

use crate::storage::{
    Snapshot, Statistics,
    metrics::CommandKind,
    mvcc::SnapshotReader,
    txn::{
        ProcessResult, RESOLVE_LOCK_BATCH_SIZE, Result,
        commands::{Command, CommandExt, ReadCommand, ResolveLock, TypedCommand},
        sched_pool::tls_collect_keyread_histogram_vec,
    },
};

command! {
    /// Scan locks for resolving according to `txn_status`.
    ///
    /// During the GC operation, this should be called to find out stale locks whose timestamp is
    /// before safe point.
    /// This should followed by a `ResolveLock`.
    ResolveLockReadPhase:
        cmd_ty => (),
        display => "kv::resolve_lock_readphase {:?} scan_key({:?}) | {:?}", (txn_status, scan_key, ctx),
        content => {
            /// Maps lock_ts to commit_ts. See ./resolve_lock.rs for details.
            txn_status: HashMap<TimeStamp, TimeStamp>,
            scan_key: Option<Key>,
            is_background: bool,
        }
}

impl CommandExt for ResolveLockReadPhase {
    ctx!();
    request_type!(KvResolveLock);
    property!(readonly);

    #[inline]
    fn tag(&self) -> CommandKind {
        if self.is_background {
            CommandKind::resolve_lock_background
        } else {
            CommandKind::resolve_lock
        }
    }

    fn write_bytes(&self) -> usize {
        0
    }

    gen_lock!(empty);
}

#[maybe_async::async_trait]
impl<S: Snapshot + 'static> ReadCommand<S> for ResolveLockReadPhase {
    #[maybe_async]
    async fn process_read(self, snapshot: S, statistics: &mut Statistics) -> Result<ProcessResult> {
        let tag = self.tag();
        let (ctx, txn_status, is_background) = (self.ctx, self.txn_status, self.is_background);
        let mut reader = SnapshotReader::new_with_ctx(TimeStamp::default(), snapshot, &ctx);
        let result = reader.scan_locks(
            self.scan_key.as_ref(),
            None,
            |lock| txn_status.contains_key(&lock.ts),
            RESOLVE_LOCK_BATCH_SIZE,
        );
        statistics.add(&reader.take_statistics());
        let (kv_pairs, has_remain) = result?;
        let mut flatten_pairs = Vec::with_capacity(kv_pairs.len());
        for (key, lock_or_shared) in kv_pairs {
            match lock_or_shared {
                Either::Left(lock) => flatten_pairs.push((key, lock)),
                Either::Right(mut shared_locks) => {
                    let ts_to_process: Vec<_> = shared_locks
                        .iter_ts()
                        .filter(|ts| txn_status.contains_key(ts))
                        .cloned()
                        .collect();
                    for ts in ts_to_process {
                        if let Some(lock) = shared_locks
                            .get_lock(&ts)
                            .map_err(crate::storage::mvcc::Error::from)
                            .map_err(crate::storage::txn::Error::from)?
                        {
                            flatten_pairs.push((key.clone(), lock.clone()));
                        }
                    }
                }
            }
        }
        tls_collect_keyread_histogram_vec(tag.get_str(), flatten_pairs.len() as f64);

        if flatten_pairs.is_empty() {
            Ok(ProcessResult::Res)
        } else {
            let next_scan_key = if has_remain {
                // There might be more locks.
                flatten_pairs.last().map(|(k, _lock)| k.clone())
            } else {
                // All locks are scanned
                None
            };
            let next_cmd = ResolveLock {
                ctx,
                deadline: self.deadline,
                txn_status,
                scan_key: next_scan_key,
                key_locks: flatten_pairs,
                txn_file_status: Default::default(),
                is_background,
            };
            Ok(ProcessResult::NextCommand {
                cmd: Command::ResolveLock(next_cmd),
            })
        }
    }
}