// Copyright 2025 TiKV Project Authors. Licensed under Apache-2.0.

use std::{sync::Arc, time::Duration};

use api_version::ApiV2;
use cloud_encryption::{
    EncryptionKey, EncryptionKeyManager, KEY_TYPE_AES_256_CTR_LEGACY, MasterKey,
    current_data_key_from_encryption_meta_bytes,
};
use dashmap::DashMap;
use encryption_conf_change::KeyspaceEncryptionMetaId;
use kvengine::{
    dfs::{Dfs, FileType, Options},
    encryption_key_from_shard_properties,
};
use load_data::dispatcher::get_shard_meta;
use pd_client::PdClient;
use tikv_util::{box_err, info, warn};

use crate::error::{Error, Result};

const GET_SHARD_META_TIMEOUT: Duration = Duration::from_secs(60);
const KEYSPACE_INFO_CACHE_TTL: Duration = Duration::from_secs(3 * 60 * 60);

#[derive(Clone, Debug)]
enum KeyspaceEncryptionInfo {
    Unencrypted,
    Encrypted(EncryptionKey),
}

#[derive(Clone, Debug)]
struct CachedKeyspaceEncryptionInfo {
    info: KeyspaceEncryptionInfo,
    cached_at: tikv_util::time::Instant,
}

#[derive(Default)]
pub(crate) struct KeyspaceEncryptionInfoCache {
    entries: DashMap<u32, CachedKeyspaceEncryptionInfo>,
}

impl KeyspaceEncryptionInfoCache {
    // NOTE: This cache is only for next-gen keyspace encryption
    // (`ENCRYPTION_META_KEY` / PD file-id based lookup). Legacy encryption keys
    // (`ENCRYPTION_KEY`) are per-shard and are not cached.
    fn get(&self, keyspace_id: u32) -> Option<KeyspaceEncryptionInfo> {
        let entry = self.entries.get(&keyspace_id)?;
        if entry.cached_at.saturating_elapsed() > KEYSPACE_INFO_CACHE_TTL {
            // Treat expired entries as cache misses; a subsequent `set` will
            // overwrite stale data.
            return None;
        }
        Some(entry.info.clone())
    }

    fn set(&self, keyspace_id: u32, info: KeyspaceEncryptionInfo) {
        self.entries.insert(
            keyspace_id,
            CachedKeyspaceEncryptionInfo {
                info,
                cached_at: tikv_util::time::Instant::now_coarse(),
            },
        );
    }
}

async fn get_encryption_meta_by_keyspace_id(
    pd_client: Arc<dyn PdClient>,
    keyspace_id: u32,
) -> Result<KeyspaceEncryptionMetaId> {
    let pd_client_for_status = pd_client;
    let join = tokio::task::spawn_blocking(move || {
        let mgr = encryption_conf_change::EncryptionConfChangeManager::new(
            pd_client_for_status,
            encryption_conf_change::OwnerCheckMode::Disabled,
        );
        mgr.load_keyspace_encryption_status_from_pd(keyspace_id)
    })
    .await;

    match join {
        Ok(Ok(status)) => Ok(status.select_meta_id(false)),
        Ok(Err(e)) if e.is_keyspace_not_found() => {
            // Some test or legacy environments do not populate keyspace meta in
            // PD. Fall back to the original shard-meta path and treat missing
            // encryption meta as unencrypted.
            Ok(KeyspaceEncryptionMetaId::UnavailableInPd)
        }
        Ok(Err(e)) => Err(Error::Other(box_err!(
            "CMEK get-meta: load_keyspace_encryption_status_from_pd failed, keyspace_id={}: {:?}",
            keyspace_id,
            e
        ))),
        Err(e) => Err(Error::Other(box_err!(
            "CMEK get-meta: load_keyspace_encryption_status_from_pd join failed, keyspace_id={}: {:?}",
            keyspace_id,
            e
        ))),
    }
}

async fn load_encryption_key_from_shard_meta(
    cache: &KeyspaceEncryptionInfoCache,
    keyspace_id: u32,
    pd_client: Arc<dyn PdClient>,
    encryption_key_manager: Arc<EncryptionKeyManager>,
    legacy_master_key: &MasterKey,
) -> Result<Option<EncryptionKey>> {
    // Fetch encryption info from TiKV's status server API.
    let (keyspace_start, keyspace_end) = ApiV2::get_txn_keyspace_range(keyspace_id);
    let shard_meta = get_shard_meta(pd_client, &keyspace_start, GET_SHARD_META_TIMEOUT)
        .await
        .map_err(|err| {
            Error::Other(box_err!(
                "get shard meta failed: {:?}, keyspace_id: {}",
                err,
                keyspace_id
            ))
        })?;
    let snapshot = shard_meta.get_snapshot();

    if snapshot.outer_start < keyspace_start || keyspace_end < snapshot.outer_end {
        // Would happen if the keyspace has not been created or has been merged.
        warn!("load_encryption_key_for_keyspace: keyspace range mismatch"; "keyspace_id" => keyspace_id, "snapshot" => ?snapshot);
        return Err(box_err!("keyspace range mismatch"));
    }

    let encryption_key = encryption_key_from_shard_properties(
        snapshot.get_properties(),
        encryption_key_manager,
        legacy_master_key,
        Some(keyspace_id),
    )?;

    match encryption_key.as_ref() {
        None => {
            cache.set(keyspace_id, KeyspaceEncryptionInfo::Unencrypted);
            info!(
                "CMEK config not found for keyspace";
                "keyspace_id" => keyspace_id,
            );
        }
        Some(key) if key.get_type() == KEY_TYPE_AES_256_CTR_LEGACY => {
            // Legacy keys are per-shard; do not cache by keyspace.
            info!(
                "CMEK legacy encryption key found for keyspace";
                "keyspace_id" => keyspace_id,
                "key_id" => ?key.get_key_id(),
            );
        }
        Some(key) => {
            info!(
                "CMEK encryption key found for keyspace";
                "keyspace_id" => keyspace_id,
                "key_id" => ?key.get_key_id(),
            );
            cache.set(keyspace_id, KeyspaceEncryptionInfo::Encrypted(key.clone()));
        }
    }

    Ok(encryption_key)
}

/// Get the encryption key for a keyspace, if encryption is enabled.
///
/// In the common case, this consults PD keyspace config and reads the stable
/// `EncryptionMeta` version from DFS using PD
/// `encryption_meta_current_file_id`, so the worker won't encrypt with a key
/// that hasn't rolled out to all regions yet.
///
/// Only when encryption is enabled but PD reports the meta id as unavailable,
/// it falls back to querying the leader store's shard meta
/// (`/kvengine/meta/<shard_id>`).
///
/// The only case when the meta id is unavailable in PD is when the keyspace is
/// created with CMEK (not through the enable CMEK operation) and it has not
/// gone through the key rotation process yet, so any region within the keyspace
/// should return the same encryption meta.
///
/// TODO: backfill PD `current_file_id` for such cases to simplify the worker
/// logic and eventually remove the shard-meta fallback.
pub(crate) async fn load_encryption_key_for_keyspace(
    cache: &KeyspaceEncryptionInfoCache,
    keyspace_id: u32,
    pd_client: Arc<dyn PdClient>,
    dfs: Arc<dyn Dfs>,
    encryption_key_manager: Arc<EncryptionKeyManager>,
    legacy_master_key: &MasterKey,
) -> Result<Option<EncryptionKey>> {
    // Fast path: keyspace-level cache only covers next-gen keyspace encryption
    // meta.
    if let Some(info) = cache.get(keyspace_id) {
        return Ok(match info {
            KeyspaceEncryptionInfo::Unencrypted => None,
            KeyspaceEncryptionInfo::Encrypted(key) => Some(key),
        });
    }

    // PD path: decide which encryption meta file ID (if any) should be used for
    // this keyspace.
    let selection = get_encryption_meta_by_keyspace_id(pd_client.clone(), keyspace_id).await?;

    match selection {
        KeyspaceEncryptionMetaId::Unencrypted => {
            cache.set(keyspace_id, KeyspaceEncryptionInfo::Unencrypted);
            info!(
                "CMEK encryption not enabled for keyspace";
                "keyspace_id" => keyspace_id,
            );
            return Ok(None);
        }
        KeyspaceEncryptionMetaId::MetaFileId(current_file_id) => {
            let opts = Options::default().with_type(FileType::EncryptionDict);
            let meta_bytes = dfs.read_file(current_file_id, opts).await?;
            let key = current_data_key_from_encryption_meta_bytes(
                encryption_key_manager,
                meta_bytes.as_ref(),
                Some(keyspace_id),
            )?;
            info!(
                "CMEK encryption key found for keyspace via PD/DFS";
                "keyspace_id" => keyspace_id,
                "key_id" => ?key.get_key_id(),
                "current_file_id" => current_file_id,
            );
            cache.set(keyspace_id, KeyspaceEncryptionInfo::Encrypted(key.clone()));
            return Ok(Some(key));
        }
        KeyspaceEncryptionMetaId::UnavailableInPd => {
            // Fall back to shard meta below.
        }
    }

    load_encryption_key_from_shard_meta(
        cache,
        keyspace_id,
        pd_client,
        encryption_key_manager,
        legacy_master_key,
    )
    .await
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_keyspace_encryption_info_cache_isolated_between_workers() {
        let first_worker_cache = KeyspaceEncryptionInfoCache::default();
        let second_worker_cache = KeyspaceEncryptionInfoCache::default();

        first_worker_cache.set(1, KeyspaceEncryptionInfo::Unencrypted);

        assert!(matches!(
            first_worker_cache.get(1),
            Some(KeyspaceEncryptionInfo::Unencrypted)
        ));
        assert!(second_worker_cache.get(1).is_none());
    }
}