use async_trait::async_trait;
use crate::error::Result;

#[async_trait]
pub trait KmsProvider: Send + Sync {
    async fn encrypt(&self, key_uri: &str, plaintext: &[u8]) -> Result<Vec<u8>>;
    async fn decrypt(&self, key_uri: &str, ciphertext: &[u8]) -> Result<Vec<u8>>;
}

/// A placeholder provider that can be used for local testing or when BYOK is disabled.
pub struct NoopKmsProvider;

#[async_trait]
impl KmsProvider for NoopKmsProvider {
    async fn encrypt(&self, _key_uri: &str, plaintext: &[u8]) -> Result<Vec<u8>> {
        if std::env::var("CANNON_ENV").unwrap_or_default() == "production" {
            return Err(crate::error::CannonError::Internal("NOOP KMS provider is not allowed in production".to_string()));
        }
        Ok(plaintext.to_vec())
    }
    async fn decrypt(&self, _key_uri: &str, ciphertext: &[u8]) -> Result<Vec<u8>> {
        if std::env::var("CANNON_ENV").unwrap_or_default() == "production" {
            return Err(crate::error::CannonError::Internal("NOOP KMS provider is not allowed in production".to_string()));
        }
        Ok(ciphertext.to_vec())
    }
}

pub struct AwsKmsProvider {
    client: aws_sdk_kms::Client,
}

impl AwsKmsProvider {
    pub async fn new() -> Self {
        let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
        let client = aws_sdk_kms::Client::new(&config);
        Self { client }
    }
}

#[async_trait]
impl KmsProvider for AwsKmsProvider {
    async fn encrypt(&self, key_uri: &str, plaintext: &[u8]) -> Result<Vec<u8>> {
        let blob = aws_sdk_kms::primitives::Blob::new(plaintext);
        let resp = self.client.encrypt()
            .key_id(key_uri)
            .plaintext(blob)
            .send()
            .await
            .map_err(|e| crate::error::CannonError::Internal(format!("AWS KMS encryption fail: {}", e)))?;
        
        Ok(resp.ciphertext_blob.map(|b| b.into_inner()).unwrap_or_default())
    }

    async fn decrypt(&self, key_uri: &str, ciphertext: &[u8]) -> Result<Vec<u8>> {
        let blob = aws_sdk_kms::primitives::Blob::new(ciphertext);
        let resp = self.client.decrypt()
            .key_id(key_uri)
            .ciphertext_blob(blob)
            .send()
            .await
            .map_err(|e| crate::error::CannonError::Internal(format!("AWS KMS decryption fail: {}", e)))?;
        
        Ok(resp.plaintext.map(|b| b.into_inner()).unwrap_or_default())
    }
}

pub struct GcpKmsProvider;
#[async_trait]
impl KmsProvider for GcpKmsProvider {
    async fn encrypt(&self, key_uri: &str, plaintext: &[u8]) -> Result<Vec<u8>> {
        if std::env::var("CANNON_ENV").unwrap_or_default() == "production" {
            return Err(crate::error::CannonError::Internal("GCP KMS provider is currently a mock and not allowed in production".to_string()));
        }
        tracing::info!("Mocking GCP KMS encryption for key {}", key_uri);
        // Future: implement google-cloud-kms logic
        Ok(plaintext.to_vec())
    }
    async fn decrypt(&self, key_uri: &str, ciphertext: &[u8]) -> Result<Vec<u8>> {
        if std::env::var("CANNON_ENV").unwrap_or_default() == "production" {
            return Err(crate::error::CannonError::Internal("GCP KMS provider is currently a mock and not allowed in production".to_string()));
        }
        tracing::info!("Mocking GCP KMS decryption for key {}", key_uri);
        Ok(ciphertext.to_vec())
    }
}

pub struct KmsFactory;
impl KmsFactory {
    pub async fn get_provider(provider_type: Option<&str>) -> std::sync::Arc<dyn KmsProvider> {
        let env = std::env::var("CANNON_ENV").unwrap_or_default();
        match provider_type {
            Some("aws") => std::sync::Arc::new(AwsKmsProvider::new().await) as std::sync::Arc<dyn KmsProvider>,
            Some("gcp") => std::sync::Arc::new(GcpKmsProvider) as std::sync::Arc<dyn KmsProvider>,
            _ => {
                if env == "production" {
                    tracing::error!("CRITICAL: KMS provider requested but none configured in production. Defaulting to NOOP (will fail on use).");
                }
                std::sync::Arc::new(NoopKmsProvider) as std::sync::Arc<dyn KmsProvider>
            }
        }
    }
}
