DualFlow/api/routers/health.py

"""
Health Check Router
"""
from fastapi import APIRouter, Depends
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
from shared.database import get_db
from shared.models.health import HealthStatus
import asyncio
import time

router = APIRouter()


@router.get("/health", response_model=HealthStatus)
async def health_check(db: Session = Depends(get_db)):
    """Health check endpoint"""
    start_time = time.time()
    
    # Check database connection
    try:
        db.execute("SELECT 1")
        db_status = "healthy"
    except Exception as e:
        db_status = f"unhealthy: {str(e)}"
    
    # Check external services
    services_status = await check_external_services()
    
    response_time = time.time() - start_time
    
    return HealthStatus(
        status="healthy" if db_status == "healthy" else "unhealthy",
        timestamp=time.time(),
        response_time=response_time,
        database=db_status,
        services=services_status
    )


@router.get("/health/ready")
async def readiness_check():
    """Readiness check for Kubernetes"""
    return {"status": "ready"}


@router.get("/health/live")
async def liveness_check():
    """Liveness check for Kubernetes"""
    return {"status": "alive"}


async def check_external_services():
    """Check external services status"""
    services = {
        "redis": await check_redis(),
        "rabbitmq": await check_rabbitmq(),
        "model_service": await check_model_service(),
        "user_service": await check_user_service(),
        "task_service": await check_task_service()
    }
    return services


async def check_redis():
    """Check Redis connection"""
    try:
        # Implementation would check Redis connection
        return "healthy"
    except Exception:
        return "unhealthy"


async def check_rabbitmq():
    """Check RabbitMQ connection"""
    try:
        # Implementation would check RabbitMQ connection
        return "healthy"
    except Exception:
        return "unhealthy"


async def check_model_service():
    """Check model service"""
    try:
        # Implementation would check model service
        return "healthy"
    except Exception:
        return "unhealthy"


async def check_user_service():
    """Check user service"""
    try:
        # Implementation would check user service
        return "healthy"
    except Exception:
        return "unhealthy"


async def check_task_service():
    """Check task service"""
    try:
        # Implementation would check task service
        return "healthy"
    except Exception:
        return "unhealthy"