"""Health check service functions for dependency monitoring."""

import asyncio
import time
from typing import Literal

import structlog
from redis.asyncio import Redis
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from src.api.schemas.health import DependencyHealth
from src.core.config import settings

logger = structlog.get_logger(__name__)

# Timeout for individual health checks (5s)
HEALTH_CHECK_TIMEOUT = 5


async def check_postgres(session: AsyncSession) -> DependencyHealth:
    """Check PostgreSQL database health.

    Executes a simple SELECT 1 query to verify database connectivity.

    Args:
        session: Async SQLAlchemy session.

    Returns:
        DependencyHealth with status and latency.
    """
    start_time = time.perf_counter()

    try:
        async with asyncio.timeout(HEALTH_CHECK_TIMEOUT):
            await session.execute(text("SELECT 1"))
            latency_ms = int((time.perf_counter() - start_time) * 1000)

            logger.debug(
                "postgres_health_check",
                status="healthy",
                latency_ms=latency_ms,
            )

            return DependencyHealth(
                status="healthy",
                latency_ms=latency_ms,
            )

    except asyncio.TimeoutError:
        latency_ms = int((time.perf_counter() - start_time) * 1000)
        error_msg = "Health check timed out"

        logger.warning(
            "postgres_health_check",
            status="unhealthy",
            error=error_msg,
            latency_ms=latency_ms,
        )

        return DependencyHealth(
            status="unhealthy",
            latency_ms=latency_ms,
            error=error_msg,
        )

    except Exception as e:
        latency_ms = int((time.perf_counter() - start_time) * 1000)
        error_msg = str(e)

        logger.warning(
            "postgres_health_check",
            status="unhealthy",
            error=error_msg,
            latency_ms=latency_ms,
        )

        return DependencyHealth(
            status="unhealthy",
            latency_ms=latency_ms,
            error=error_msg,
        )


async def check_redis(redis_client: Redis) -> DependencyHealth:
    """Check Redis health.

    Executes a PING command to verify Redis connectivity.

    Args:
        redis_client: Async Redis client.

    Returns:
        DependencyHealth with status and latency.
    """
    start_time = time.perf_counter()

    try:
        async with asyncio.timeout(HEALTH_CHECK_TIMEOUT):
            response = await redis_client.ping()
            latency_ms = int((time.perf_counter() - start_time) * 1000)

            if response:
                logger.debug(
                    "redis_health_check",
                    status="healthy",
                    latency_ms=latency_ms,
                )

                return DependencyHealth(
                    status="healthy",
                    latency_ms=latency_ms,
                )
            else:
                error_msg = "PING did not return expected response"

                logger.warning(
                    "redis_health_check",
                    status="unhealthy",
                    error=error_msg,
                    latency_ms=latency_ms,
                )

                return DependencyHealth(
                    status="unhealthy",
                    latency_ms=latency_ms,
                    error=error_msg,
                )

    except asyncio.TimeoutError:
        latency_ms = int((time.perf_counter() - start_time) * 1000)
        error_msg = "Health check timed out"

        logger.warning(
            "redis_health_check",
            status="unhealthy",
            error=error_msg,
            latency_ms=latency_ms,
        )

        return DependencyHealth(
            status="unhealthy",
            latency_ms=latency_ms,
            error=error_msg,
        )

    except Exception as e:
        latency_ms = int((time.perf_counter() - start_time) * 1000)
        error_msg = str(e)

        logger.warning(
            "redis_health_check",
            status="unhealthy",
            error=error_msg,
            latency_ms=latency_ms,
        )

        return DependencyHealth(
            status="unhealthy",
            latency_ms=latency_ms,
            error=error_msg,
        )


async def check_milvus() -> DependencyHealth:
    """Check Milvus vector store health.

    Connects to Milvus and lists collections to verify connectivity.

    Returns:
        DependencyHealth with status and latency.
    """
    start_time = time.perf_counter()

    try:
        async with asyncio.timeout(HEALTH_CHECK_TIMEOUT):
            from pymilvus import connections, utility

            connections.connect(
                alias="health_check",
                host=settings.milvus_host,
                port=settings.milvus_port,
            )

            # List collections to verify connection
            utility.list_collections(using="health_check")

            connections.disconnect("health_check")

            latency_ms = int((time.perf_counter() - start_time) * 1000)

            logger.debug(
                "milvus_health_check",
                status="healthy",
                latency_ms=latency_ms,
            )

            return DependencyHealth(
                status="healthy",
                latency_ms=latency_ms,
            )

    except asyncio.TimeoutError:
        latency_ms = int((time.perf_counter() - start_time) * 1000)
        error_msg = "Health check timed out"

        logger.warning(
            "milvus_health_check",
            status="unhealthy",
            error=error_msg,
            latency_ms=latency_ms,
        )

        return DependencyHealth(
            status="unhealthy",
            latency_ms=latency_ms,
            error=error_msg,
        )

    except Exception as e:
        latency_ms = int((time.perf_counter() - start_time) * 1000)
        error_msg = str(e)

        logger.warning(
            "milvus_health_check",
            status="unhealthy",
            error=error_msg,
            latency_ms=latency_ms,
        )

        return DependencyHealth(
            status="unhealthy",
            latency_ms=latency_ms,
            error=error_msg,
        )


def aggregate_health_status(
    dependencies: dict[str, DependencyHealth],
) -> Literal["healthy", "degraded", "unhealthy"]:
    """Aggregate individual dependency health into overall status.

    Args:
        dependencies: Dict of dependency name to health status.

    Returns:
        Overall status: healthy, degraded, or unhealthy.
    """
    all_healthy = all(dep.status == "healthy" for dep in dependencies.values())

    if all_healthy:
        return "healthy"

    # Check if database is unhealthy (critical dependency)
    db_status = dependencies.get("database")
    if db_status and db_status.status == "unhealthy":
        return "unhealthy"

    # Non-critical dependencies unhealthy = degraded
    return "degraded"