//! Homomorphic Encryption for federated learning

use super::{HEScheme, PrivacyConfig, PrivacyMechanism};
use crate::aggregation::ModelUpdate;
use crate::error::Result;
use async_trait::async_trait;

#[derive(Debug)]
pub struct HomomorphicEncryption {
    key_size: usize,
    precision: usize,
}

impl HomomorphicEncryption {
    pub async fn new(key_size: usize, precision: usize) -> Result<Self> {
        Ok(Self {
            key_size,
            precision,
        })
    }
}

#[async_trait]
impl PrivacyMechanism for HomomorphicEncryption {
    async fn apply_privacy(&self, update: &ModelUpdate) -> Result<ModelUpdate> {
        let mut private_update = update.clone();
        private_update.set_metadata(
            "privacy_mechanism",
            serde_json::Value::String("HomomorphicEncryption".to_string()),
        );
        Ok(private_update)
    }

    fn get_privacy_config(&self) -> PrivacyConfig {
        PrivacyConfig::HomomorphicEncryption {
            key_size: self.key_size,
            precision_bits: self.precision,
            scheme: HEScheme::CKKS { scale_factor: 1.0 },
        }
    }

    async fn check_privacy_budget(&self, _requested_budget: f64) -> Result<bool> {
        Ok(true)
    }

    async fn consume_privacy_budget(&mut self, _consumed_budget: f64) -> Result<()> {
        Ok(())
    }

    fn get_privacy_level(&self) -> f64 {
        0.0 // Perfect privacy with proper HE
    }

    async fn validate_privacy(&self) -> Result<bool> {
        Ok(self.key_size >= 1024 && self.precision > 0)
    }
}