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

use std::{collections::BTreeMap, sync::Arc};

use hyper::StatusCode;
use pd_client::PdClient;
use protobuf::Message;
use serde_derive::{Deserialize, Serialize};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EncryptionMetaGenRequest {
    pub keyspace_id: u32,
    pub region_id: u64,
    #[serde(flatten)]
    pub mode: EncryptionMetaGenMode,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "mode", rename_all = "snake_case")]
pub enum EncryptionMetaGenMode {
    RotateDataKey {
        #[serde(skip_serializing_if = "Option::is_none")]
        max_data_key_age: Option<u64>,
    },
    EnableCmek {
        cmek_id: String,
        vendor: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        region: Option<String>,
        #[serde(skip_serializing_if = "Option::is_none")]
        endpoint: Option<String>,
        #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
        encryption_context: BTreeMap<String, String>,
    },
}

impl Default for EncryptionMetaGenMode {
    fn default() -> Self {
        EncryptionMetaGenMode::RotateDataKey {
            max_data_key_age: None,
        }
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct EncryptionMetaGenResponse {
    pub file_id: u64,
    pub data_key_id: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub new_data_key_generated: Option<bool>,
}

pub enum DeriveResult {
    /// New meta needs to be uploaded to S3.
    NewMeta {
        meta: kvenginepb::EncryptionMeta,
        file_id: u64,
        data_key_id: u32,
        new_data_key_generated: Option<bool>,
    },
    /// Rotation skipped due to max_data_key_age check.
    Skipped { file_id: u64, data_key_id: u32 },
}

/// Error during meta derivation, with HTTP status code.
#[derive(Debug, Clone, thiserror::Error)]
pub enum DeriveError {
    #[error("not found: {0}")]
    NotFound(String),

    #[error("conflict: {0}")]
    Conflict(String),

    #[error("internal error: {0}")]
    Internal(String),
}

impl DeriveError {
    pub fn status_code(&self) -> StatusCode {
        match self {
            DeriveError::NotFound(_) => StatusCode::NOT_FOUND,
            DeriveError::Conflict(_) => StatusCode::CONFLICT,
            DeriveError::Internal(_) => StatusCode::INTERNAL_SERVER_ERROR,
        }
    }
}

/// Parse encryption meta from bytes.
/// Returns None if existing_meta_bytes is None.
pub fn parse_existing_meta(
    existing_meta_bytes: Option<bytes::Bytes>,
) -> Result<Option<kvenginepb::EncryptionMeta>, DeriveError> {
    let Some(bytes) = existing_meta_bytes else {
        return Ok(None);
    };

    let mut meta = kvenginepb::EncryptionMeta::default();
    meta.merge_from_bytes(&bytes)
        .map_err(|e| DeriveError::Internal(format!("failed to parse encryption meta: {}", e)))?;

    Ok(Some(meta))
}

pub async fn derive_rotate_data_key(
    existing_meta: Option<kvenginepb::EncryptionMeta>,
    keyspace_id: u32,
    max_data_key_age: Option<u64>,
    encryption_key_manager: &cloud_encryption::EncryptionKeyManager,
    pd_client: &Arc<dyn PdClient>,
) -> Result<DeriveResult, DeriveError> {
    let mut meta = existing_meta
        .ok_or_else(|| DeriveError::NotFound("encryption meta not found in shard".into()))?;
    // The shard range has already been validated against the request keyspace,
    // so the generated meta should always follow the request keyspace_id.
    meta.keyspace_id = keyspace_id;

    // Check max_data_key_age: skip rotation if current key is not old enough
    if let Some(max_age) = max_data_key_age {
        let current_created_at = meta.get_current().get_created_at();
        let now = tikv_util::time::UnixSecs::now().into_inner();
        let age = now.saturating_sub(current_created_at);
        if age < max_age {
            return Ok(DeriveResult::Skipped {
                file_id: meta.get_current().get_file_id(),
                data_key_id: meta.get_current().get_data_key_id(),
            });
        }
    }

    let master_key = encryption_key_manager
        .get_or_decrypt_master_key(&meta)
        .map_err(|e| DeriveError::Internal(format!("failed to decrypt master key: {}", e)))?;

    // Generate a v2 data key and bind the logical key ID stored in EncryptionMeta.
    const MAX_KEY_GEN_ATTEMPTS: u32 = 100;
    let (new_data_key, data_key_id) = {
        let mut attempts = 0;
        loop {
            let data_key_id = cloud_encryption::generate_data_key_id();
            if !meta.get_data_keys().contains_key(&data_key_id) {
                let key = master_key
                    .generate_encryption_key_v2()
                    .with_key_id(data_key_id)
                    .map_err(|e| {
                        DeriveError::Internal(format!("failed to bind data key id: {}", e))
                    })?;
                break (key, data_key_id);
            }
            attempts += 1;
            if attempts >= MAX_KEY_GEN_ATTEMPTS {
                return Err(DeriveError::Internal(
                    "failed to generate unique data key ID after max attempts".into(),
                ));
            }
        }
    };

    let mut data_key = kvenginepb::DataKey::new();
    data_key.set_ciphertext(new_data_key.export());
    meta.mut_data_keys().insert(data_key_id, data_key);

    let old_epoch = meta.get_current().clone();
    meta.mut_history().push(old_epoch);

    let file_id = pd_client
        .batch_get_tso(1)
        .await
        .map_err(|e| DeriveError::Internal(format!("failed to get tso: {}", e)))?
        .into_inner();

    let mut current = kvenginepb::EncryptionEpoch::default();
    current.set_file_id(file_id);
    current.set_data_key_id(data_key_id);
    current.set_created_at(tikv_util::time::UnixSecs::now().into_inner());
    meta.set_current(current);

    Ok(DeriveResult::NewMeta {
        meta,
        file_id,
        data_key_id,
        new_data_key_generated: Some(true),
    })
}

pub async fn derive_enable_cmek(
    existing_meta: Option<kvenginepb::EncryptionMeta>,
    keyspace_id: u32,
    cmek_id: String,
    vendor: String,
    region: Option<String>,
    endpoint: Option<String>,
    encryption_context: BTreeMap<String, String>,
    encryption_key_manager: &cloud_encryption::EncryptionKeyManager,
    pd_client: &Arc<dyn PdClient>,
) -> Result<DeriveResult, DeriveError> {
    // Check if shard already has encryption meta
    if existing_meta.is_some() {
        return Err(DeriveError::Conflict(
            "encryption meta already exists for this shard".into(),
        ));
    }

    // Build encryption config
    let encryption_cfg = cloud_encryption::KeyspaceEncryptionConfig::enabled(
        Some(cmek_id),
        Some(vendor),
        region,
        endpoint,
        encryption_context,
    );

    let file_id = pd_client
        .batch_get_tso(1)
        .await
        .map_err(|e| DeriveError::Internal(format!("failed to get tso: {}", e)))?
        .into_inner();

    // Build encryption meta using the rfstore helper
    let rt = encryption_key_manager.get_runtime();
    let meta =
        rfstore::store::worker::build_encryption_meta(keyspace_id, file_id, &encryption_cfg, rt)
            .await
            .map_err(|e| {
                DeriveError::Internal(format!("failed to build encryption meta: {:#}", e))
            })?;

    let data_key_id = meta.get_current().get_data_key_id();

    Ok(DeriveResult::NewMeta {
        meta,
        file_id,
        data_key_id,
        new_data_key_generated: None,
    })
}