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:

    def __init__(self):
        self.services: Dict[str, ServiceConfig] = {}
        self.service_health: Dict[str, ServiceHealth] = {}
        self.session: Optional[aiohttp.ClientSession] = None

    async def initialize(self):
        self.session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=10)
        )
        logger.info("Service registry initialized")

    async def close(self):
        if self.session:
            await self.session.close()
        logger.info("Service registry closed")

    async def register_services(self, services: Dict[str, ServiceConfig]):
        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):
        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):
        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:
        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):
        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)

            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 {self.service_health.values()}")

    async def get_service_health(self, service_name: str) -> Optional[ServiceHealth]:
        return self.service_health.get(service_name)

    async def get_all_service_health(self) -> Dict[str, Dict]:
        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:
        health = self.service_health.get(service_name)
        return health is not None and health.status == "healthy"

    async def get_healthy_services(self) -> List[str]:
        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]:
        return self.services.get(service_name)

    def get_all_services(self) -> Dict[str, ServiceConfig]:
        return self.services.copy()

    async def get_service_url(self, service_name: str) -> Optional[str]:
        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