//! FedNova implementation with normalized averaging for non-IID data

use super::{AggregationAlgorithm, AggregationConfig, ModelUpdate};
use crate::error::{FederatedError, Result};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tracing::debug;
use uuid::Uuid;

/// FedNova aggregation algorithm for normalized averaging
#[derive(Debug)]
pub struct FedNova {
    config: FedNovaConfig,
    tau_effective: f64,
    round_count: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FedNovaConfig {
    pub tau_effective: f64,
    pub momentum_factor: f64,
    pub normalize_weights: bool,
}

impl FedNova {
    pub async fn new(tau_eff: f64, momentum: f64) -> Result<Self> {
        let config = FedNovaConfig {
            tau_effective: tau_eff,
            momentum_factor: momentum,
            normalize_weights: true,
        };

        Ok(Self {
            config,
            tau_effective: tau_eff,
            round_count: 0,
        })
    }
}

#[async_trait]
impl AggregationAlgorithm for FedNova {
    async fn aggregate(&self, updates: &[ModelUpdate]) -> Result<ModelUpdate> {
        if updates.is_empty() {
            return Err(FederatedError::NoUpdatesCollected);
        }

        debug!(
            "🔄 Starting FedNova aggregation with {} updates",
            updates.len()
        );

        // Compute normalized weights based on local epochs
        let mut weights = Vec::new();
        for update in updates {
            let local_steps = update.local_epochs as f64;
            let normalized_weight = local_steps / self.tau_effective;
            weights.push(normalized_weight * update.sample_count as f64);
        }

        let aggregated_params = super::WeightedAggregator::weighted_average(updates, &weights)?;

        let mut aggregated_update = ModelUpdate::new(Uuid::new_v4());
        aggregated_update.parameters = aggregated_params;

        let total_samples: usize = updates.iter().map(|u| u.sample_count).sum();
        aggregated_update.sample_count = total_samples;
        aggregated_update.set_metadata(
            "aggregation_algorithm",
            serde_json::Value::String("FedNova".to_string()),
        );

        Ok(aggregated_update)
    }

    fn get_config(&self) -> AggregationConfig {
        AggregationConfig::FedNova {
            tau_effective: self.config.tau_effective,
            momentum_factor: self.config.momentum_factor,
            normalize_weights: self.config.normalize_weights,
        }
    }

    async fn update_state(&mut self, round: usize, _updates: &[ModelUpdate]) -> Result<()> {
        self.round_count = round;
        Ok(())
    }

    fn supports_async(&self) -> bool {
        false
    }

    fn name(&self) -> &'static str {
        "FedNova"
    }
}