// Copyright 2025 PingCAP, Inc.

use std::sync::Arc;

use quick_cache::{Weighter, sync::GuardResult};
use tikv_util::sys::SysQuota;

use crate::{
    metrics::{
        ENGINE_COLUMNAR_META_CACHE_ENTRIES, ENGINE_COLUMNAR_META_CACHE_HIT,
        ENGINE_COLUMNAR_META_CACHE_MISS, ENGINE_COLUMNAR_META_CACHE_PARSE,
        ENGINE_COLUMNAR_META_CACHE_REPORTED_SIZE_BYTES, ENGINE_COLUMNAR_META_CACHE_WEIGHTED_SIZE,
    },
    table::{Result, columnar::TableMeta},
};

const TABLE_META_CACHE_SHARDS_PER_CORE: usize = 16;
const ESTIMATED_TABLE_META_SIZE: u64 = 1024; // 1KB

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ColumnarMetaCacheKey {
    pub file_id: u64,
    pub table_id: i64,
}

#[derive(Clone)]
pub struct ColumnarMetaWeighter;

impl quick_cache::Weighter<ColumnarMetaCacheKey, Arc<TableMeta>> for ColumnarMetaWeighter {
    fn weight(&self, _: &ColumnarMetaCacheKey, val: &Arc<TableMeta>) -> u64 {
        8 + 8 + val.size
    }
}

/// Accumulated eviction deltas collected within a single cache request.
#[derive(Default)]
struct EvictionDelta {
    /// Number of evicted entries.
    entries: u64,
    /// Total weight of evicted entries.
    weight: u64,
}

/// A [`quick_cache::Lifecycle`] that keeps the entries and weighted_size gauges
/// in sync with the cache.
#[derive(Clone, Default)]
struct ColumnarMetaLifecycle;

impl quick_cache::Lifecycle<ColumnarMetaCacheKey, Arc<TableMeta>> for ColumnarMetaLifecycle {
    type RequestState = EvictionDelta;

    fn begin_request(&self) -> Self::RequestState {
        EvictionDelta::default()
    }

    fn on_evict_cold(
        &self,
        delta: &mut Self::RequestState,
        key: ColumnarMetaCacheKey,
        val: Arc<TableMeta>,
    ) {
        delta.entries += 1;
        delta.weight += ColumnarMetaWeighter.weight(&key, &val);
    }

    fn on_evict_hot(
        &self,
        delta: &mut Self::RequestState,
        key: ColumnarMetaCacheKey,
        val: Arc<TableMeta>,
    ) {
        delta.entries += 1;
        delta.weight += ColumnarMetaWeighter.weight(&key, &val);
    }

    fn end_request(&self, delta: Self::RequestState) {
        // Called after the shard lock is released; gauge updates are atomic.
        if delta.entries > 0 {
            ENGINE_COLUMNAR_META_CACHE_ENTRIES.sub(delta.entries as i64);
            ENGINE_COLUMNAR_META_CACHE_WEIGHTED_SIZE.sub(delta.weight as i64);
        }
    }
}

#[derive(Clone)]
pub struct ColumnarMetaCache {
    cache: Arc<
        quick_cache::sync::Cache<
            ColumnarMetaCacheKey,
            Arc<TableMeta>,
            ColumnarMetaWeighter,
            quick_cache::DefaultHashBuilder,
            ColumnarMetaLifecycle,
        >,
    >,
}

impl Default for ColumnarMetaCache {
    fn default() -> Self {
        let mem_limit = SysQuota::memory_limit_in_bytes();
        Self::new(mem_limit / 500)
    }
}

impl ColumnarMetaCache {
    pub fn new(max_capacity: u64) -> Self {
        let columnar_meta_cache_shards =
            (SysQuota::cpu_cores_quota() as usize).max(1) * TABLE_META_CACHE_SHARDS_PER_CORE;
        let estimated_items = (max_capacity / ESTIMATED_TABLE_META_SIZE).max(1) as usize;
        let opts = quick_cache::OptionsBuilder::new()
            .shards(columnar_meta_cache_shards)
            .weight_capacity(max_capacity)
            .estimated_items_capacity(estimated_items)
            .build()
            .unwrap();
        let cache = quick_cache::sync::Cache::with_options(
            opts,
            ColumnarMetaWeighter,
            quick_cache::DefaultHashBuilder::default(),
            ColumnarMetaLifecycle,
        );
        Self {
            cache: Arc::new(cache),
        }
    }

    pub(crate) fn try_get_with(
        &self,
        file_id: u64,
        table_id: i64,
        init: impl FnOnce() -> Result<Arc<TableMeta>>,
    ) -> Result<Arc<TableMeta>> {
        let key = ColumnarMetaCacheKey { file_id, table_id };
        match self.cache.get_value_or_guard(&key, None) {
            GuardResult::Value(v) => {
                ENGINE_COLUMNAR_META_CACHE_HIT.inc();
                Ok(v)
            }
            GuardResult::Guard(guard) => {
                ENGINE_COLUMNAR_META_CACHE_MISS.inc();
                let meta = init()?;
                // Count PARSE / reported size only after a successful fill. IO
                // failures in `init` must not look like TableMeta::parse work.
                ENGINE_COLUMNAR_META_CACHE_PARSE.inc();
                ENGINE_COLUMNAR_META_CACHE_REPORTED_SIZE_BYTES.inc_by(meta.size);
                // `guard.insert` runs the lifecycle internally: evictions are
                // reflected in the gauges via `end_request` (outside the shard
                // lock). The lifecycle does NOT observe the insertion itself,
                // so bump the gauges only when insert returns Ok.
                //
                // Overweight reject path: quick_cache calls `on_evict_cold` for
                // the never-admitted value (−) and still returns Ok; the manual
                // +(weight) below cancels that out so gauges stay consistent.
                if guard.insert(meta.clone()).is_ok() {
                    ENGINE_COLUMNAR_META_CACHE_ENTRIES.inc();
                    ENGINE_COLUMNAR_META_CACHE_WEIGHTED_SIZE
                        .add(ColumnarMetaWeighter.weight(&key, &meta) as i64);
                }
                Ok(meta)
            }
            GuardResult::Timeout => unreachable!(),
        }
    }
}