ForcePilot/backend/package/yuxi/channel/interfaces/rest/router/health.py

from __future__ import annotations

import asyncio
import logging
import os

from fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlalchemy import select

from yuxi.channel.container import ChannelContainer, get_channel

logger = logging.getLogger(__name__)

router = APIRouter()


class HealthzResponse(BaseModel):
    status: str = "ok"
    startup_timeline: list[dict] | None = None


class ReadyzComponentStatus(BaseModel):
    status: str
    detail: str | None = None


class ReadyzResponse(BaseModel):
    status: str
    redis: ReadyzComponentStatus | None = None
    postgres: ReadyzComponentStatus | None = None
    workers: ReadyzComponentStatus | None = None
    channels: ReadyzComponentStatus | None = None
    ws_connections: ReadyzComponentStatus | None = None


@router.get("/healthz", response_model=HealthzResponse)
async def liveness_check(channel: ChannelContainer = Depends(get_channel)):
    show_timeline = os.getenv("CHANNEL_SHOW_STARTUP_TIMELINE", "false").lower() == "true"
    return HealthzResponse(
        status="ok",
        startup_timeline=channel.startup_tracer.to_dict() if show_timeline else None,
    )


@router.get("/readyz")
async def readiness_check(channel: ChannelContainer = Depends(get_channel)):
    keys = ["redis", "postgres", "workers", "channels", "ws_connections"]
    results = await asyncio.gather(
        _check_redis(channel),
        _check_postgres(),
        _check_workers(channel),
        _check_channels(channel),
        _check_ws_connections(channel),
    )
    checks = dict(zip(keys, results))

    has_error = any(v.status == "error" for v in checks.values())
    has_degraded = any(v.status == "degraded" for v in checks.values())

    if has_error:
        overall, status_code = "not_ready", 503
    elif has_degraded:
        overall, status_code = "degraded", 200
    else:
        overall, status_code = "ready", 200

    return JSONResponse(
        status_code=status_code,
        content=ReadyzResponse(status=overall, **checks).model_dump(),
    )


async def _check_redis(container: ChannelContainer) -> ReadyzComponentStatus:
    try:
        if container and container.redis:
            await container.redis.ping()
            return ReadyzComponentStatus(status="ok")
    except Exception as exc:
        logger.warning("redis check failed: %s", exc)
    return ReadyzComponentStatus(status="error", detail="connection refused")


async def _check_postgres() -> ReadyzComponentStatus:
    try:
        from yuxi.storage.postgres.manager import pg_manager

        if not pg_manager.is_initialized:
            return ReadyzComponentStatus(status="error", detail="not initialized")
        async with pg_manager.get_async_session_context() as session:
            await session.execute(select(1))
        return ReadyzComponentStatus(status="ok")
    except Exception as exc:
        logger.warning("postgres check failed: %s", exc)
    return ReadyzComponentStatus(status="error", detail="connection refused")


async def _check_workers(container: ChannelContainer) -> ReadyzComponentStatus:
    try:
        if container and container.worker_pool and container.worker_pool.is_running:
            return ReadyzComponentStatus(status="ok")
    except Exception as exc:
        logger.warning("workers check failed: %s", exc)
    return ReadyzComponentStatus(status="error", detail="no active consumers")


async def _check_channels(container: ChannelContainer) -> ReadyzComponentStatus:
    try:
        if container and container.adapters:
            results = await asyncio.gather(
                *[a.is_healthy() for a in container.adapters.values()],
                return_exceptions=True,
            )
            unhealthy = [name for name, r in zip(container.adapters, results) if r is not True]
            if not unhealthy:
                return ReadyzComponentStatus(status="ok")
            return ReadyzComponentStatus(
                status="degraded",
                detail=f"unhealthy channels: {', '.join(unhealthy)}",
            )
    except Exception as exc:
        logger.warning("channels check failed: %s", exc)
    return ReadyzComponentStatus(status="error", detail="no adapters")


async def _check_ws_connections(container: ChannelContainer) -> ReadyzComponentStatus:
    if not container or not container.ws_manager:
        return ReadyzComponentStatus(status="ok", detail="no ws connections")
    status_map = container.ws_manager.connections_status
    if not status_map:
        return ReadyzComponentStatus(status="ok", detail="no ws connections registered")
    disconnected = [ct for ct, ok in status_map.items() if not ok]
    if not disconnected:
        return ReadyzComponentStatus(status="ok")
    return ReadyzComponentStatus(
        status="degraded",
        detail=f"disconnected: {', '.join(disconnected)}",
    )