"""
Service registry for managing microservice discovery and health monitoring
"""

import aiohttp
import asyncio
from datetime import datetime
from typing import Dict, List, Optional
import logging

from models import ServiceConfig, ServiceHealth

logger = logging.getLogger(__name__)

class ServiceRegistry:
    """Service registry for microservice management"""
    
    def __init__(self):
        self.services: Dict[str, ServiceConfig] = {}
        self.service_health: Dict[str, ServiceHealth] = {}
        self.session: Optional[aiohttp.ClientSession] = None
    
    async def initialize(self):
        """Initialize the service registry"""
        self.session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=10)
        )
        logger.info("Service registry initialized")
    
    async def close(self):
        """Close the service registry"""
        if self.session:
            await self.session.close()
        logger.info("Service registry closed")
    
    async def register_services(self, services: Dict[str, ServiceConfig]):
        """Register multiple services"""
        self.services.update(services)
        
        # Initialize health status for all services
        for service_name, config in services.items():
            self.service_health[service_name] = ServiceHealth(
                service=service_name,
                status="unknown",
                last_check=datetime.utcnow()
            )
        
        logger.info(f"Registered {len(services)} services")
        
        # Perform initial health check
        await self.update_all_service_health()
    
    async def register_service(self, service_config: ServiceConfig):
        """Register a single service"""
        self.services[service_config.name] = service_config
        self.service_health[service_config.name] = ServiceHealth(
            service=service_config.name,
            status="unknown",
            last_check=datetime.utcnow()
        )
        
        logger.info(f"Registered service: {service_config.name}")
        
        # Check health of the newly registered service
        await self.check_service_health(service_config.name)
    
    async def unregister_service(self, service_name: str):
        """Unregister a service"""
        self.services.pop(service_name, None)
        self.service_health.pop(service_name, None)
        logger.info(f"Unregistered service: {service_name}")
    
    async def check_service_health(self, service_name: str) -> ServiceHealth:
        """Check health of a specific service"""
        service_config = self.services.get(service_name)
        if not service_config:
            logger.error(f"Service {service_name} not found in registry")
            return ServiceHealth(
                service=service_name,
                status="unknown",
                last_check=datetime.utcnow(),
                error_message="Service not registered"
            )
        
        start_time = datetime.utcnow()
        
        try:
            health_url = f"{service_config.base_url}{service_config.health_endpoint}"
            
            async with self.session.get(health_url) as response:
                end_time = datetime.utcnow()
                response_time = (end_time - start_time).total_seconds() * 1000
                
                if response.status == 200:
                    health_data = await response.json()
                    status = "healthy" if health_data.get("status") in ["healthy", "ok"] else "unhealthy"
                    
                    health = ServiceHealth(
                        service=service_name,
                        status=status,
                        response_time_ms=response_time,
                        last_check=end_time
                    )
                else:
                    health = ServiceHealth(
                        service=service_name,
                        status="unhealthy",
                        response_time_ms=response_time,
                        last_check=end_time,
                        error_message=f"HTTP {response.status}"
                    )
                
        except aiohttp.ClientError as e:
            health = ServiceHealth(
                service=service_name,
                status="unhealthy",
                last_check=datetime.utcnow(),
                error_message=f"Connection error: {str(e)}"
            )
        except Exception as e:
            health = ServiceHealth(
                service=service_name,
                status="unhealthy",
                last_check=datetime.utcnow(),
                error_message=f"Health check failed: {str(e)}"
            )
        
        # Update health status
        self.service_health[service_name] = health
        
        # Log health status changes
        if health.status != "healthy":
            logger.warning(f"Service {service_name} health check failed: {health.error_message}")
        
        return health
    
    async def update_all_service_health(self):
        """Update health status for all registered services"""
        health_checks = [
            self.check_service_health(service_name)
            for service_name in self.services.keys()
        ]
        
        if health_checks:
            await asyncio.gather(*health_checks, return_exceptions=True)
            
            # Log summary
            healthy_count = sum(1 for h in self.service_health.values() if h.status == "healthy")
            total_count = len(self.services)
            logger.info(f"Health check complete: {healthy_count}/{total_count} services healthy")
    
    async def get_service_health(self, service_name: str) -> Optional[ServiceHealth]:
        """Get health status of a specific service"""
        return self.service_health.get(service_name)
    
    async def get_all_service_health(self) -> Dict[str, Dict]:
        """Get health status of all services"""
        health_dict = {}
        for service_name, health in self.service_health.items():
            health_dict[service_name] = {
                "status": health.status,
                "response_time_ms": health.response_time_ms,
                "last_check": health.last_check.isoformat(),
                "error_message": health.error_message
            }
        return health_dict
    
    async def is_service_healthy(self, service_name: str) -> bool:
        """Check if a service is healthy"""
        health = self.service_health.get(service_name)
        return health is not None and health.status == "healthy"
    
    async def get_healthy_services(self) -> List[str]:
        """Get list of healthy service names"""
        return [
            service_name
            for service_name, health in self.service_health.items()
            if health.status == "healthy"
        ]
    
    def get_service_config(self, service_name: str) -> Optional[ServiceConfig]:
        """Get configuration for a specific service"""
        return self.services.get(service_name)
    
    def get_all_services(self) -> Dict[str, ServiceConfig]:
        """Get all registered services"""
        return self.services.copy()
    
    async def get_service_url(self, service_name: str) -> Optional[str]:
        """Get base URL for a healthy service"""
        if await self.is_service_healthy(service_name):
            service_config = self.services.get(service_name)
            return service_config.base_url if service_config else None
        return None