"""
Load Balancer for Toxo System.

Manages load balancing for distributed operations and resource allocation.
"""

import asyncio
from typing import Dict, Any, List, Optional
from datetime import datetime
import random

from ..utils.logger import get_logger


class LoadBalancer:
    """Basic load balancer for Toxo system."""
    
    def __init__(self):
        self.logger = get_logger("toxo.core.load_balancer")
        self.nodes: List[Dict[str, Any]] = []
        self.current_index = 0
        self.logger.info("Load balancer initialized")
    
    def add_node(self, node_id: str, capacity: int = 100) -> None:
        """Add a node to the load balancer."""
        node = {
            "id": node_id,
            "capacity": capacity,
            "current_load": 0,
            "status": "active",
            "last_used": datetime.now()
        }
        self.nodes.append(node)
        self.logger.info(f"Added node {node_id} with capacity {capacity}")
    
    def get_best_node(self) -> Optional[Dict[str, Any]]:
        """Get the best available node based on load."""
        if not self.nodes:
            return None
        
        # Find node with lowest load
        active_nodes = [node for node in self.nodes if node["status"] == "active"]
        if not active_nodes:
            return None
        
        best_node = min(active_nodes, key=lambda n: n["current_load"] / n["capacity"])
        return best_node
    
    def assign_task(self, task_weight: int = 1) -> Optional[str]:
        """Assign a task to the best available node."""
        node = self.get_best_node()
        if node and node["current_load"] + task_weight <= node["capacity"]:
            node["current_load"] += task_weight
            node["last_used"] = datetime.now()
            self.logger.debug(f"Assigned task (weight {task_weight}) to node {node['id']}")
            return node["id"]
        return None
    
    def release_task(self, node_id: str, task_weight: int = 1) -> None:
        """Release a task from a node."""
        for node in self.nodes:
            if node["id"] == node_id:
                node["current_load"] = max(0, node["current_load"] - task_weight)
                self.logger.debug(f"Released task (weight {task_weight}) from node {node_id}")
                break
    
    def get_status(self) -> Dict[str, Any]:
        """Get load balancer status."""
        total_capacity = sum(node["capacity"] for node in self.nodes)
        total_load = sum(node["current_load"] for node in self.nodes)
        
        return {
            "total_nodes": len(self.nodes),
            "active_nodes": len([n for n in self.nodes if n["status"] == "active"]),
            "total_capacity": total_capacity,
            "total_load": total_load,
            "utilization": total_load / max(1, total_capacity),
            "nodes": self.nodes.copy()
        }
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "total_nodes": len(getattr(self, 'nodes', [])),
            "active_nodes": len([n for n in getattr(self, 'nodes', []) if n.get("status") == "active"]),
            "system_type": "LoadBalancer"
        } 