"""
Health check registry and management for LaneSwap services.

This module provides the health check decorator and registry system
for custom health checks in LaneSwap-managed services.
"""

import asyncio
import inspect
import time
from typing import Any, Callable, Dict, List, Optional, Union

HealthCheckFunc = Callable[[], Union[Dict[str, Any], None]]
AsyncHealthCheckFunc = Callable[[], Union[Dict[str, Any], None]]


class HealthCheckRegistry:
    """
    Registry for health check functions.

    Manages registration and execution of custom health checks,
    supporting both synchronous and asynchronous functions.
    """

    def __init__(self) -> None:
        """Initialize empty health check registry."""
        self._checks: List[HealthCheckFunc] = []
        self._async_checks: List[AsyncHealthCheckFunc] = []

    def register(self, func: Union[HealthCheckFunc, AsyncHealthCheckFunc]) -> None:
        """
        Register a health check function.

        Args:
            func: Health check function (sync or async)

        Examples:
            >>> registry = HealthCheckRegistry()
            >>> def check_db():
            ...     return {"database": "connected"}
            >>> registry.register(check_db)
        """
        if inspect.iscoroutinefunction(func):
            self._async_checks.append(func)
        else:
            self._checks.append(func)

    def run_checks(self) -> Dict[str, Any]:
        """
        Run all registered synchronous health checks.

        Returns:
            Dictionary containing results from all health checks

        Examples:
            >>> registry = Health CheckRegistry()
            >>> registry.register(lambda: {"service": "ok"})
            >>> registry.run_checks()
            {'service': 'ok'}
        """
        results: Dict[str, Any] = {}

        for check in self._checks:
            try:
                result = check()
                if result is not None and isinstance(result, dict):
                    results.update(result)
            except Exception as e:
                # Health check failed - log error but don't crash
                results["error"] = str(e)

        return results

    async def run_async_checks(self) -> Dict[str, Any]:
        """
        Run all registered asynchronous health checks.

        Returns:
            Dictionary containing results from all async health checks

        Examples:
            >>> registry = HealthCheckRegistry()
            >>> async def check_redis():
            ...     return {"redis": "connected"}
            >>> registry.register(check_redis)
            >>> await registry.run_async_checks()
            {'redis': 'connected'}
        """
        results: Dict[str, Any] = {}

        for check in self._async_checks:
            try:
                result = await check()
                if result is not None and isinstance(result, dict):
                    results.update(result)
            except Exception as e:
                # Health check failed - log error but don't crash
                results["error"] = str(e)

        return results

    def has_async_checks(self) -> bool:
        """
        Check if any async health checks are registered.

        Returns:
            True if async checks exist, False otherwise
        """
        return len(self._async_checks) > 0

    def clear(self) -> None:
        """Clear all registered health checks."""
        self._checks.clear()
        self._async_checks.clear()


def build_health_response(
    service_name: str,
    custom_checks: Optional[Dict[str, Any]] = None,
    status: str = "healthy",
) -> Dict[str, Any]:
    """
    Build standard LaneSwap health check response.

    Args:
        service_name: Name of the service
        custom_checks: Optional custom health check results
        status: Health status (healthy, degraded, unhealthy)

    Returns:
        Dictionary formatted for LaneSwap Monitor

    Examples:
        >>> build_health_response("my-service")
        {
            'status': 'healthy',
            'service': 'my-service',
            'timestamp': 1234567890.123
        }

        >>> build_health_response("my-service", {"database": "connected"})
        {
            'status': 'healthy',
            'service': 'my-service',
            'timestamp': 1234567890.123,
            'database': 'connected'
        }
    """
    response: Dict[str, Any] = {
        "status": status,
        "service": service_name,
        "timestamp": time.time(),
    }

    # Merge custom health check results
    if custom_checks:
        response.update(custom_checks)

    return response
