import logging
import random
from typing import List, Dict

logger = logging.getLogger(__name__)

class LoadBalancer:
    def __init__(self):
        self.workers: Dict[str, float] = {}  # worker_id -> performance_score
        self.worker_load: Dict[str, int] = {}  # worker_id -> current_load
        self.worker_health: Dict[str, bool] = {}  # worker_id -> is_healthy
        
    def add_worker(self, worker_id: str, performance_score: float):
        """Aggiunge un worker al load balancer"""
        self.workers[worker_id] = performance_score
        self.worker_load[worker_id] = 0
        self.worker_health[worker_id] = True
        logger.info(f"✅ Added worker {worker_id} with score {performance_score}")
        
    def remove_worker(self, worker_id: str):
        """Rimuove un worker dal load balancer"""
        if worker_id in self.workers:
            del self.workers[worker_id]
        if worker_id in self.worker_load:
            del self.worker_load[worker_id]
        if worker_id in self.worker_health:
            del self.worker_health[worker_id]
        logger.info(f"🗑️ Removed worker {worker_id}")
        
    def update_worker_score(self, worker_id: str, new_score: float):
        """NUOVO: Aggiorna il punteggio di performance di un worker"""
        if worker_id in self.workers:
            old_score = self.workers[worker_id]
            self.workers[worker_id] = new_score
            logger.info(f"📊 Updated worker {worker_id} score: {old_score:.2f} -> {new_score:.2f}")
        else:
            logger.warning(f"⚠️ Cannot update score: worker {worker_id} not found")
            
    def set_worker_health(self, worker_id: str, is_healthy: bool):
        """NUOVO: Imposta lo stato di salute di un worker"""
        if worker_id in self.worker_health:
            self.worker_health[worker_id] = is_healthy
            status = "healthy" if is_healthy else "unhealthy"
            logger.info(f"❤️ Worker {worker_id} marked as {status}")
        
    def get_healthy_workers(self) -> List[str]:
        """NUOVO: Restituisce solo i worker healthy"""
        return [worker_id for worker_id in self.workers 
                if self.worker_health.get(worker_id, True)]
        
    def distribute_work(self, num_tasks: int) -> List[str]:
        """Distribuisce il lavoro tra i worker disponibili"""
        healthy_workers = self.get_healthy_workers()
        
        if not healthy_workers:
            logger.warning("⚠️ No healthy workers available for distribution")
            return []
            
        # Ordina worker healthy per performance score (decrescente)
        sorted_workers = sorted(
            [(wid, self.workers[wid]) for wid in healthy_workers], 
            key=lambda x: x[1], 
            reverse=True
        )
        
        # Distribuisce i task ai worker migliori
        assigned_workers = []
        available_workers = [worker_id for worker_id, _ in sorted_workers]
        
        for i in range(num_tasks):
            if available_workers:
                worker_id = available_workers[i % len(available_workers)]
                assigned_workers.append(worker_id)
                self.worker_load[worker_id] += 1
                
        logger.info(f"📦 Distributed {num_tasks} tasks to {len(assigned_workers)} workers")
        return assigned_workers
        
    def get_worker_load(self, worker_id: str) -> int:
        """Ottiene il carico corrente di un worker"""
        return self.worker_load.get(worker_id, 0)
        
    def complete_task(self, worker_id: str):
        """Segna un task come completato per un worker"""
        if worker_id in self.worker_load and self.worker_load[worker_id] > 0:
            self.worker_load[worker_id] -= 1
            logger.debug(f"✅ Task completed by worker {worker_id}")
            
    def get_best_worker(self) -> str:
        """NUOVO: Restituisce il worker con il punteggio più alto"""
        if not self.workers:
            return None
            
        healthy_workers = self.get_healthy_workers()
        if not healthy_workers:
            return None
            
        return max(healthy_workers, key=lambda wid: self.workers[wid])
            
    def get_total_performance_score(self) -> float:
        """Restituisce il punteggio prestazionale totale"""
        healthy_workers = self.get_healthy_workers()
        return sum(self.workers[wid] for wid in healthy_workers) if healthy_workers else 0.0
            
    def get_worker_stats(self) -> Dict:
        """Ottiene statistiche sui worker"""
        healthy_workers = self.get_healthy_workers()
        
        return {
            'total_workers': len(self.workers),
            'healthy_workers': len(healthy_workers),
            'unhealthy_workers': len(self.workers) - len(healthy_workers),
            'worker_loads': self.worker_load.copy(),
            'performance_scores': self.workers.copy(),
            'health_status': self.worker_health.copy(),
            'total_performance': self.get_total_performance_score(),
            'best_worker': self.get_best_worker()
        }
        
    def redistribute_load(self):
        """NUOVO: Ribilancia il carico tra i worker"""
        total_load = sum(self.worker_load.values())
        healthy_workers = self.get_healthy_workers()
        
        if not healthy_workers:
            return
            
        avg_load = total_load // len(healthy_workers)
        
        for worker_id in healthy_workers:
            self.worker_load[worker_id] = avg_load
            
        logger.info(f"⚖️ Load redistributed: {avg_load} tasks per worker")