"""
Load balancer for distributing requests across service instances
"""

import random
from typing import List, Dict, Optional
import logging

logger = logging.getLogger(__name__)

class LoadBalancer:
    """Simple load balancer for microservice requests"""

    def __init__(self, service_registry=None):
        # In a real implementation, this would track multiple instances per service
        self.service_instances: Dict[str, List[str]] = {}
        self.current_index: Dict[str, int] = {}
        self.service_registry = service_registry
    
    def register_service_instance(self, service_name: str, instance_url: str):
        """Register a new service instance"""
        if service_name not in self.service_instances:
            self.service_instances[service_name] = []
            self.current_index[service_name] = 0
        
        if instance_url not in self.service_instances[service_name]:
            self.service_instances[service_name].append(instance_url)
            logger.info(f"Registered instance {instance_url} for service {service_name}")
    
    def unregister_service_instance(self, service_name: str, instance_url: str):
        """Unregister a service instance"""
        if service_name in self.service_instances:
            try:
                self.service_instances[service_name].remove(instance_url)
                logger.info(f"Unregistered instance {instance_url} for service {service_name}")
                
                # Reset index if it's out of bounds
                if self.current_index[service_name] >= len(self.service_instances[service_name]):
                    self.current_index[service_name] = 0
                    
            except ValueError:
                logger.warning(f"Instance {instance_url} not found for service {service_name}")
    
    async def get_service_url(self, service_name: str, strategy: str = "single") -> Optional[str]:
        """
        Get a service URL using the specified load balancing strategy
        
        Strategies:
        - single: Single instance (default for this simple implementation)
        - round_robin: Round-robin across instances
        - random: Random selection
        """
        # For this microservice setup, we typically have one instance per service
        # In a production environment, you'd have multiple instances
        
        if strategy == "single":
            # Default behavior - get the service URL from service registry
            if self.service_registry:
                return await self.service_registry.get_service_url(service_name)
            else:
                logger.error("No service registry available")
                return None
        
        elif strategy == "round_robin":
            return await self._round_robin_select(service_name)
        
        elif strategy == "random":
            return await self._random_select(service_name)
        
        else:
            logger.error(f"Unknown load balancing strategy: {strategy}")
            return None
    
    async def _round_robin_select(self, service_name: str) -> Optional[str]:
        """Select service instance using round-robin"""
        instances = self.service_instances.get(service_name, [])
        if not instances:
            # Fall back to service registry
            if self.service_registry:
                return await self.service_registry.get_service_url(service_name)
            else:
                logger.error("No service registry available for fallback")
                return None
        
        # Round-robin selection
        current_idx = self.current_index[service_name]
        selected_instance = instances[current_idx]
        
        # Update index for next request
        self.current_index[service_name] = (current_idx + 1) % len(instances)
        
        logger.debug(f"Round-robin selected {selected_instance} for {service_name}")
        return selected_instance
    
    async def _random_select(self, service_name: str) -> Optional[str]:
        """Select service instance randomly"""
        instances = self.service_instances.get(service_name, [])
        if not instances:
            # Fall back to service registry
            if self.service_registry:
                return await self.service_registry.get_service_url(service_name)
            else:
                logger.error("No service registry available for fallback")
                return None
        
        selected_instance = random.choice(instances)
        logger.debug(f"Random selected {selected_instance} for {service_name}")
        return selected_instance
    
    def get_service_instances(self, service_name: str) -> List[str]:
        """Get all registered instances for a service"""
        return self.service_instances.get(service_name, [])
    
    def get_instance_count(self, service_name: str) -> int:
        """Get number of registered instances for a service"""
        return len(self.service_instances.get(service_name, []))
    
    def get_all_services(self) -> Dict[str, List[str]]:
        """Get all services and their instances"""
        return self.service_instances.copy()
    
    def health_check_failed(self, service_name: str, instance_url: str):
        """Handle health check failure for a service instance"""
        logger.warning(f"Health check failed for {instance_url} ({service_name})")
        # In a production system, you might temporarily remove unhealthy instances
        # For now, we just log the failure
    
    def health_check_recovered(self, service_name: str, instance_url: str):
        """Handle health check recovery for a service instance"""
        logger.info(f"Health check recovered for {instance_url} ({service_name})")
        # Re-register the instance if it was temporarily removed