//! Reputation-based Byzantine protection

use super::{ByzantineConfig, ByzantineRobust};
use crate::aggregation::ModelUpdate;
use crate::error::{FederatedError, Result};
use async_trait::async_trait;
use std::collections::HashMap;
use uuid::Uuid;

#[derive(Debug)]
pub struct ReputationSystem {
    reputation_scores: HashMap<Uuid, f64>,
    initial_reputation: f64,
    decay_factor: f64,
    min_reputation: f64,
}

impl ReputationSystem {
    pub async fn new(initial_reputation: f64, min_reputation: f64) -> Result<Self> {
        Ok(Self {
            reputation_scores: HashMap::new(),
            initial_reputation,
            decay_factor: 0.95,
            min_reputation,
        })
    }
}

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

        // Filter based on reputation scores
        let filtered: Vec<ModelUpdate> = updates
            .iter()
            .filter(|update| {
                let reputation = self
                    .reputation_scores
                    .get(&update.client_id)
                    .unwrap_or(&self.initial_reputation);
                *reputation >= self.min_reputation
            })
            .cloned()
            .collect();

        Ok(filtered)
    }

    async fn detect_malicious_clients(&self, updates: &[ModelUpdate]) -> Result<Vec<Uuid>> {
        let malicious: Vec<Uuid> = updates
            .iter()
            .filter_map(|update| {
                let reputation = self
                    .reputation_scores
                    .get(&update.client_id)
                    .unwrap_or(&self.initial_reputation);
                if *reputation < self.min_reputation {
                    Some(update.client_id)
                } else {
                    None
                }
            })
            .collect();

        Ok(malicious)
    }

    fn get_robustness_config(&self) -> ByzantineConfig {
        ByzantineConfig::ReputationBased {
            initial_reputation: self.initial_reputation,
            decay_factor: self.decay_factor,
            min_reputation: self.min_reputation,
        }
    }

    async fn update_state(&mut self, _round: usize, updates: &[ModelUpdate]) -> Result<()> {
        // Update reputation scores based on behavior
        for update in updates {
            let current_reputation = self
                .reputation_scores
                .entry(update.client_id)
                .or_insert(self.initial_reputation);

            // Simple reputation update based on loss improvement
            if update.loss < 1.0 {
                // Good behavior
                *current_reputation = (*current_reputation * 0.9 + 1.0 * 0.1).min(1.0);
            } else {
                // Poor behavior
                *current_reputation = (*current_reputation * 0.9).max(0.0);
            }
        }
        Ok(())
    }

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