"""Federated learning infrastructure for privacy-preserving AI improvement."""

from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
import asyncio
import numpy as np
from abc import ABC, abstractmethod
# import torch  # Optional - using numpy fallback
import hashlib
import json

from ..utils.logging import get_logger


logger = get_logger(__name__)


@dataclass
class FederatedClient:
    """Represents a federated learning client."""
    client_id: str
    last_update: datetime
    model_version: str
    contribution_count: int = 0
    trust_score: float = 1.0
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ModelUpdate:
    """Model update from a federated client."""
    client_id: str
    timestamp: datetime
    gradients: Optional[Dict[str, np.ndarray]] = None
    metrics: Dict[str, float] = field(default_factory=dict)
    data_size: int = 0
    update_hash: str = ""


@dataclass
class FederatedRound:
    """A round of federated learning."""
    round_id: int
    start_time: datetime
    end_time: Optional[datetime] = None
    participating_clients: List[str] = field(default_factory=list)
    aggregated_update: Optional[Dict[str, np.ndarray]] = None
    metrics: Dict[str, float] = field(default_factory=dict)


class PrivacyMechanism(ABC):
    """Abstract base class for privacy mechanisms."""
    
    @abstractmethod
    async def apply_privacy(self, data: np.ndarray) -> np.ndarray:
        """Apply privacy mechanism to data."""
        pass


class DifferentialPrivacy(PrivacyMechanism):
    """Differential privacy implementation."""
    
    def __init__(self, epsilon: float = 1.0, delta: float = 1e-5):
        self.epsilon = epsilon
        self.delta = delta
    
    async def apply_privacy(self, data: np.ndarray) -> np.ndarray:
        """Add noise for differential privacy."""
        # Calculate sensitivity (simplified)
        sensitivity = np.max(np.abs(data)) * 2
        
        # Calculate noise scale
        noise_scale = sensitivity / self.epsilon
        
        # Add Laplacian noise
        noise = np.random.laplace(0, noise_scale, data.shape)
        
        return data + noise
    
    def privacy_budget_remaining(self, queries_used: int) -> float:
        """Calculate remaining privacy budget."""
        return max(0, self.epsilon - (queries_used * 0.1))


class SecureAggregation:
    """Secure aggregation for federated learning."""
    
    def __init__(self, threshold: int = 3):
        self.threshold = threshold  # Minimum clients for aggregation
    
    async def aggregate_updates(
        self,
        updates: List[ModelUpdate],
        weights: Optional[List[float]] = None
    ) -> Dict[str, np.ndarray]:
        """Securely aggregate model updates."""
        if len(updates) < self.threshold:
            raise ValueError(f"Need at least {self.threshold} clients for secure aggregation")
        
        if weights is None:
            # Weight by data size
            total_data = sum(u.data_size for u in updates)
            weights = [u.data_size / total_data for u in updates]
        
        # Initialize aggregated gradients
        aggregated = {}
        
        for i, update in enumerate(updates):
            if update.gradients is None:
                continue
            
            for param_name, gradient in update.gradients.items():
                if param_name not in aggregated:
                    aggregated[param_name] = np.zeros_like(gradient)
                
                # Weighted aggregation
                aggregated[param_name] += weights[i] * gradient
        
        return aggregated
    
    def verify_update(self, update: ModelUpdate) -> bool:
        """Verify the integrity of an update."""
        # Verify hash
        computed_hash = self._compute_update_hash(update)
        return computed_hash == update.update_hash
    
    def _compute_update_hash(self, update: ModelUpdate) -> str:
        """Compute hash of model update."""
        hasher = hashlib.sha256()
        
        # Hash gradients
        if update.gradients:
            for param_name in sorted(update.gradients.keys()):
                hasher.update(param_name.encode())
                hasher.update(update.gradients[param_name].tobytes())
        
        # Hash metadata
        hasher.update(update.client_id.encode())
        hasher.update(str(update.data_size).encode())
        
        return hasher.hexdigest()


class FederatedLearningServer:
    """Central server for federated learning coordination."""
    
    def __init__(
        self,
        min_clients: int = 5,
        rounds_per_epoch: int = 10,
        privacy_mechanism: Optional[PrivacyMechanism] = None
    ):
        self.min_clients = min_clients
        self.rounds_per_epoch = rounds_per_epoch
        self.privacy_mechanism = privacy_mechanism or DifferentialPrivacy()
        
        self.clients: Dict[str, FederatedClient] = {}
        self.current_round: Optional[FederatedRound] = None
        self.round_history: List[FederatedRound] = []
        self.global_model_version = "v0.1.0"
        
        self.aggregator = SecureAggregation()
    
    async def register_client(self, client_id: str, metadata: Dict[str, Any] = None) -> bool:
        """Register a new federated learning client."""
        if client_id in self.clients:
            logger.warning(f"Client {client_id} already registered")
            return False
        
        client = FederatedClient(
            client_id=client_id,
            last_update=datetime.utcnow(),
            model_version=self.global_model_version,
            metadata=metadata or {}
        )
        
        self.clients[client_id] = client
        logger.info(f"Registered new federated client: {client_id}")
        
        return True
    
    async def start_round(self) -> FederatedRound:
        """Start a new federated learning round."""
        if self.current_round and self.current_round.end_time is None:
            raise RuntimeError("Previous round not completed")
        
        # Select participating clients
        eligible_clients = [
            c for c in self.clients.values()
            if c.trust_score > 0.5 and c.model_version == self.global_model_version
        ]
        
        if len(eligible_clients) < self.min_clients:
            raise ValueError(f"Not enough eligible clients: {len(eligible_clients)}/{self.min_clients}")
        
        # Create new round
        round_id = len(self.round_history) + 1
        self.current_round = FederatedRound(
            round_id=round_id,
            start_time=datetime.utcnow(),
            participating_clients=[c.client_id for c in eligible_clients[:self.min_clients * 2]]
        )
        
        logger.info(f"Started federated round {round_id} with {len(self.current_round.participating_clients)} clients")
        
        return self.current_round
    
    async def submit_update(self, update: ModelUpdate) -> bool:
        """Submit a model update from a client."""
        if not self.current_round:
            logger.error("No active federated round")
            return False
        
        if update.client_id not in self.current_round.participating_clients:
            logger.error(f"Client {update.client_id} not participating in current round")
            return False
        
        # Verify update integrity
        if not self.aggregator.verify_update(update):
            logger.error(f"Update from {update.client_id} failed integrity check")
            self._penalize_client(update.client_id)
            return False
        
        # Apply privacy mechanism
        if update.gradients and self.privacy_mechanism:
            for param_name, gradient in update.gradients.items():
                update.gradients[param_name] = await self.privacy_mechanism.apply_privacy(gradient)
        
        # Store update (in production, use persistent storage)
        # For now, we'll process immediately
        
        # Update client stats
        client = self.clients.get(update.client_id)
        if client:
            client.last_update = datetime.utcnow()
            client.contribution_count += 1
        
        return True
    
    async def aggregate_round(self, updates: List[ModelUpdate]) -> Dict[str, np.ndarray]:
        """Aggregate updates for the current round."""
        if not self.current_round:
            raise RuntimeError("No active round")
        
        # Filter updates from participating clients
        valid_updates = [
            u for u in updates
            if u.client_id in self.current_round.participating_clients
        ]
        
        if len(valid_updates) < self.min_clients:
            logger.warning(f"Insufficient updates: {len(valid_updates)}/{self.min_clients}")
            return {}
        
        # Perform secure aggregation
        aggregated = await self.aggregator.aggregate_updates(valid_updates)
        
        # Store results
        self.current_round.aggregated_update = aggregated
        self.current_round.end_time = datetime.utcnow()
        
        # Calculate metrics
        self.current_round.metrics = {
            "num_updates": len(valid_updates),
            "avg_data_size": np.mean([u.data_size for u in valid_updates]),
            "round_duration": (self.current_round.end_time - self.current_round.start_time).total_seconds()
        }
        
        # Move to history
        self.round_history.append(self.current_round)
        self.current_round = None
        
        # Update global model version
        self._increment_model_version()
        
        return aggregated
    
    def get_client_stats(self, client_id: str) -> Optional[Dict[str, Any]]:
        """Get statistics for a specific client."""
        client = self.clients.get(client_id)
        if not client:
            return None
        
        return {
            "client_id": client.client_id,
            "contribution_count": client.contribution_count,
            "trust_score": client.trust_score,
            "last_update": client.last_update.isoformat(),
            "model_version": client.model_version
        }
    
    def get_global_stats(self) -> Dict[str, Any]:
        """Get global federated learning statistics."""
        return {
            "total_clients": len(self.clients),
            "active_clients": len([c for c in self.clients.values() if c.trust_score > 0.5]),
            "total_rounds": len(self.round_history),
            "current_model_version": self.global_model_version,
            "avg_round_duration": np.mean([
                (r.end_time - r.start_time).total_seconds()
                for r in self.round_history
                if r.end_time
            ]) if self.round_history else 0
        }
    
    def _penalize_client(self, client_id: str) -> None:
        """Penalize a client for bad behavior."""
        client = self.clients.get(client_id)
        if client:
            client.trust_score *= 0.9
            logger.warning(f"Penalized client {client_id}, trust score: {client.trust_score}")
    
    def _increment_model_version(self) -> None:
        """Increment the global model version."""
        parts = self.global_model_version.split('.')
        parts[-1] = str(int(parts[-1]) + 1)
        self.global_model_version = '.'.join(parts)


class FederatedLearningClient:
    """Client-side federated learning implementation."""
    
    def __init__(
        self,
        client_id: str,
        local_data_size: int,
        privacy_budget: float = 10.0
    ):
        self.client_id = client_id
        self.local_data_size = local_data_size
        self.privacy_budget = privacy_budget
        self.queries_used = 0
    
    async def compute_update(
        self,
        model: Any,
        local_data: Any,
        epochs: int = 1
    ) -> ModelUpdate:
        """Compute model update on local data."""
        # Simulate local training (in production, actual training happens here)
        start_time = time.time()
        
        # Placeholder for gradient computation
        gradients = {}
        
        # In production, this would involve:
        # 1. Loading local data
        # 2. Running forward/backward passes
        # 3. Computing gradients
        # 4. Applying differential privacy
        
        # For demo, create synthetic gradients
        import time
        for param_name in ["layer1.weight", "layer2.weight", "output.weight"]:
            gradients[param_name] = np.random.randn(100, 100) * 0.01
        
        # Track privacy budget
        self.queries_used += epochs
        
        # Create update
        update = ModelUpdate(
            client_id=self.client_id,
            timestamp=datetime.utcnow(),
            gradients=gradients,
            metrics={
                "local_loss": np.random.random(),
                "local_accuracy": 0.85 + np.random.random() * 0.1
            },
            data_size=self.local_data_size
        )
        
        # Add hash for integrity
        update.update_hash = self._compute_hash(update)
        
        logger.info(f"Client {self.client_id} computed update in {time.time() - start_time:.2f}s")
        
        return update
    
    def privacy_budget_remaining(self) -> float:
        """Check remaining privacy budget."""
        return max(0, self.privacy_budget - self.queries_used)
    
    def _compute_hash(self, update: ModelUpdate) -> str:
        """Compute hash for update integrity."""
        hasher = hashlib.sha256()
        
        if update.gradients:
            for param_name in sorted(update.gradients.keys()):
                hasher.update(param_name.encode())
                hasher.update(update.gradients[param_name].tobytes())
        
        hasher.update(update.client_id.encode())
        hasher.update(str(update.data_size).encode())
        
        return hasher.hexdigest()


class LoveFederatedLearning:
    """Federated learning with love-based principles."""
    
    def __init__(self):
        self.love_metrics = {
            "collaboration_score": 0.0,
            "mutual_benefit": 0.0,
            "privacy_respect": 0.0,
            "inclusive_participation": 0.0
        }
    
    async def evaluate_federation_health(
        self,
        server: FederatedLearningServer
    ) -> Dict[str, Any]:
        """Evaluate the health of the federation with love metrics."""
        stats = server.get_global_stats()
        
        # Calculate love-based metrics
        self.love_metrics["collaboration_score"] = min(
            1.0,
            stats["active_clients"] / max(stats["total_clients"], 1)
        )
        
        self.love_metrics["inclusive_participation"] = 1.0 - (
            np.std([c.contribution_count for c in server.clients.values()]) /
            max(np.mean([c.contribution_count for c in server.clients.values()]), 1)
        )
        
        self.love_metrics["privacy_respect"] = min(
            1.0,
            sum(c.trust_score for c in server.clients.values()) / max(len(server.clients), 1)
        )
        
        # Overall health score
        health_score = np.mean(list(self.love_metrics.values()))
        
        return {
            "health_score": health_score,
            "love_metrics": self.love_metrics,
            "recommendations": self._get_recommendations(health_score)
        }
    
    def _get_recommendations(self, health_score: float) -> List[str]:
        """Get recommendations for improving federation health."""
        recommendations = []
        
        if health_score < 0.7:
            recommendations.append("Increase client engagement through better incentives")
        
        if self.love_metrics["privacy_respect"] < 0.8:
            recommendations.append("Strengthen privacy guarantees to build trust")
        
        if self.love_metrics["inclusive_participation"] < 0.7:
            recommendations.append("Ensure fair participation opportunities for all clients")
        
        return recommendations