// Copyright 2025 PingCAP, Inc.

use std::{
    fmt::{Debug, Formatter},
    hash::Hash,
    sync::Arc,
};

use bytes::Bytes;
use quick_cache::{Equivalent, sync::GuardResult};

use crate::{
    SnapAccess, UserMeta,
    metrics::{ENGINE_VALUE_CACHE_CACHE_FILL, ENGINE_VALUE_CACHE_CACHE_HIT},
    table::{AtomicSnapVersion, InnerKey, OwnedInnerKey, SnapVersion},
};

// ValueCache caches the values in the engine.
// When the reader performs a get operation, it first checks the cache, if the
// cache exists and valid, it uses the cached value directly.
#[derive(Clone)]
pub struct ValueCache {
    inner: Arc<quick_cache::sync::Cache<ValueCacheKey, ValueCacheValue, Weighter>>,
    double_check: bool,
}

impl ValueCache {
    pub fn is_double_check(&self) -> bool {
        self.double_check
    }
}

impl Debug for ValueCache {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "ValueCache")
    }
}

#[derive(Copy, Clone)]
struct Weighter;

impl quick_cache::Weighter<ValueCacheKey, ValueCacheValue> for Weighter {
    fn weight(&self, key: &ValueCacheKey, value: &ValueCacheValue) -> u64 {
        // The weight is the size of the key and value.
        key.inner_key.len() as u64 + value.value.len() as u64 + 64 // 64 bytes for metadata
    }
}

impl ValueCache {
    pub fn new(capacity: u64, double_check: bool) -> Self {
        let estimated_items_capacity = (capacity / 128) as usize;
        let inner = Arc::new(quick_cache::sync::Cache::with_weighter(
            estimated_items_capacity,
            capacity,
            Weighter,
        ));
        ValueCache {
            inner,
            double_check,
        }
    }

    pub fn get(
        &self,
        snap_access: &SnapAccess,
        outer_key: &[u8],
        start_ts: u64,
        has_lock: bool,
    ) -> Option<ValueCacheValue> {
        let keyspace_id = snap_access.get_keyspace_id();
        let inner_key = InnerKey::from_outer_key(outer_key);
        let key_ref = ValueCacheKeyRef {
            shard_ver: snap_access.get_version() as u32,
            keyspace_id,
            inner_key,
        };
        let val = self.inner.get(&key_ref)?;
        if val.version > start_ts {
            // The start_ts is too old, we cannot use the cached value.
            return None;
        }
        let val_snap_version = val.valid_snap_version.load();
        if val_snap_version >= snap_access.get_mem_table_snap_version() {
            // The shard is not changed, we can use the cached value.
            ENGINE_VALUE_CACHE_CACHE_HIT.inc();
            return Some(val);
        }
        // The shard is changed, we need to check if the value is still valid.
        let mut owned_val: Vec<u8> = vec![];
        let Some(newer_val) =
            snap_access.validate_cached_value(inner_key, val_snap_version, &mut owned_val)
        else {
            // The value may have newer version in L1+, remove it to reduce future
            // validation cost.
            let _ = self.inner.remove(&key_ref);
            return None;
        };
        if !newer_val.is_valid() {
            // The value is not changed, we update the valid_snap_version to reduce future
            // validation cost.
            if !has_lock {
                let _ = val
                    .valid_snap_version
                    .compare_exchange(val_snap_version, snap_access.get_mem_table_snap_version());
            }
            ENGINE_VALUE_CACHE_CACHE_HIT.inc();
            return Some(val);
        }
        // L0 values are generated by flush doesn't contains tombstone.
        debug_assert!(newer_val.user_meta_len() > 0);
        let um = UserMeta::from_slice(newer_val.user_meta());
        let new_start_ts = um.start_ts;
        let cache_value = ValueCacheValue::new(
            Bytes::copy_from_slice(newer_val.get_value()),
            new_start_ts,
            newer_val.version,
            snap_access.get_mem_table_snap_version(),
        );
        if !has_lock {
            ENGINE_VALUE_CACHE_CACHE_FILL.inc();
            self.set(snap_access, outer_key, cache_value.clone());
        }
        if start_ts >= newer_val.version {
            ENGINE_VALUE_CACHE_CACHE_HIT.inc();
            return Some(cache_value);
        }
        None
    }

    // set is called by the reader after get the row from the Shard.
    pub fn set(&self, snap_access: &SnapAccess, outer_key: &[u8], value: ValueCacheValue) {
        let inner_key = InnerKey::from_outer_key(outer_key);
        let cache_key = ValueCacheKey {
            shard_ver: snap_access.get_version() as u32,
            keyspace_id: snap_access.get_keyspace_id(),
            inner_key: inner_key.into(),
        };
        match self.inner.get_value_or_guard(&cache_key, None) {
            GuardResult::Value(old_val) => {
                if old_val.version < value.version {
                    self.inner.insert(cache_key.clone(), value);
                }
            }
            GuardResult::Guard(g) => {
                let _ = g.insert(value);
            }
            GuardResult::Timeout => {}
        }
    }
}

#[derive(Eq, PartialEq, Clone, Hash)]
pub struct ValueCacheKey {
    pub shard_ver: u32,
    pub keyspace_id: u32,
    pub inner_key: OwnedInnerKey,
}

impl Debug for ValueCacheKey {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "ValueCacheKey(shard_ver: {}, keyspace_id: {}, inner_key: {:?})",
            self.shard_ver, self.keyspace_id, self.inner_key,
        )
    }
}

#[derive(Eq, PartialEq, Clone, Hash)]
pub struct ValueCacheKeyRef<'a> {
    pub shard_ver: u32,
    pub keyspace_id: u32,
    pub inner_key: InnerKey<'a>,
}

impl Equivalent<ValueCacheKey> for ValueCacheKeyRef<'_> {
    fn equivalent(&self, key: &ValueCacheKey) -> bool {
        key.shard_ver == self.shard_ver
            && key.keyspace_id == self.keyspace_id
            && key.inner_key.as_ref() == self.inner_key
    }
}

impl Debug for ValueCacheKeyRef<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "ValueCacheKeyRef(shard_ver: {}, keyspace_id: {}, inner_key: {:?})",
            self.shard_ver, self.keyspace_id, self.inner_key,
        )
    }
}

#[derive(Clone)]
pub struct ValueCacheValue {
    pub value: Bytes,
    pub start_ts: u64,
    pub version: u64,
    pub(crate) valid_snap_version: Arc<AtomicSnapVersion>,
}

impl Debug for ValueCacheValue {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "ValueCacheValue(val_len: {}, start_ts: {}, version: {}, shard_write_seq: {})",
            self.value.len(),
            self.start_ts,
            self.version,
            self.valid_snap_version.load(),
        )
    }
}

impl ValueCacheValue {
    pub fn new(value: Bytes, start_ts: u64, version: u64, shard_snap_version: SnapVersion) -> Self {
        ValueCacheValue {
            value,
            start_ts,
            version,
            valid_snap_version: Arc::new(AtomicSnapVersion::new(shard_snap_version)),
        }
    }
}