# mcpgateway/main.py
@app.get("/ready")
async def ready():
    # Only checks if app is running, not pool health
    return {"status": "ready"}

# mcpgateway/metrics/pool_metrics.py

from prometheus_client import Gauge, Counter

# Pool utilization gauges
db_pool_size = Gauge(
    "db_pool_size",
    "Total size of the database connection pool"
)
db_pool_checked_out = Gauge(
    "db_pool_checked_out",
    "Number of connections currently checked out from pool"
)
db_pool_overflow = Gauge(
    "db_pool_overflow",
    "Number of overflow connections currently in use"
)
db_pool_checked_in = Gauge(
    "db_pool_checked_in",
    "Number of connections available in pool"
)
db_pool_utilization = Gauge(
    "db_pool_utilization",
    "Pool utilization ratio (0.0-1.0)"
)

# Pool event counters
db_pool_timeout_total = Counter(
    "db_pool_timeout_total",
    "Total number of pool timeout errors"
)
db_pool_checkout_total = Counter(
    "db_pool_checkout_total",
    "Total connection checkouts from pool"
)
db_pool_invalidated_total = Counter(
    "db_pool_invalidated_total",
    "Total connections invalidated"
)


def update_pool_metrics(engine):
    """Update pool metrics from SQLAlchemy engine."""
    pool = engine.pool
    
    size = pool.size()
    checked_out = pool.checkedout()
    overflow = pool.overflow()
    checked_in = pool.checkedin()
    
    db_pool_size.set(size)
    db_pool_checked_out.set(checked_out)
    db_pool_overflow.set(overflow)
    db_pool_checked_in.set(checked_in)
    
    # Utilization = checked_out / (size + max_overflow)
    max_capacity = size + pool._max_overflow
    utilization = checked_out / max_capacity if max_capacity > 0 else 0
    db_pool_utilization.set(utilization)

# mcpgateway/db.py - Add event listeners

from sqlalchemy import event

@event.listens_for(engine, "checkout")
def on_checkout(dbapi_conn, connection_record, connection_proxy):
    """Track connection checkouts."""
    db_pool_checkout_total.inc()
    update_pool_metrics(engine)

@event.listens_for(engine, "checkin")
def on_checkin(dbapi_conn, connection_record):
    """Track connection returns."""
    update_pool_metrics(engine)

@event.listens_for(engine, "invalidate")
def on_invalidate(dbapi_conn, connection_record, exception):
    """Track connection invalidations."""
    db_pool_invalidated_total.inc()
    update_pool_metrics(engine)

# mcpgateway/main.py

from mcpgateway.config import settings

@app.get("/ready")
async def ready():
    """Readiness check including connection pool health.
    
    Fails readiness when pool utilization exceeds threshold,
    causing load balancer to stop routing traffic to this replica.
    """
    pool = engine.pool
    
    # Calculate pool utilization
    checked_out = pool.checkedout()
    max_capacity = pool.size() + pool._max_overflow
    utilization = checked_out / max_capacity if max_capacity > 0 else 0
    
    pool_healthy = utilization < settings.readiness_pool_threshold
    
    # Optional: Test actual connectivity
    db_connectable = True
    if settings.readiness_check_db_connection:
        try:
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
        except Exception:
            db_connectable = False
    
    is_ready = pool_healthy and db_connectable
    
    status = {
        "status": "ready" if is_ready else "not_ready",
        "pool": {
            "utilization": round(utilization, 2),
            "checked_out": checked_out,
            "capacity": max_capacity,
            "healthy": pool_healthy,
        },
        "database": {
            "connectable": db_connectable,
        }
    }
    
    if not is_ready:
        return JSONResponse(status_code=503, content=status)
    
    return status

@app.get("/health")
async def health():
    """Liveness check with optional pool details."""
    pool = engine.pool
    
    checked_out = pool.checkedout()
    max_capacity = pool.size() + pool._max_overflow
    utilization = checked_out / max_capacity if max_capacity > 0 else 0
    
    return {
        "status": "healthy",
        "pool": {
            "size": pool.size(),
            "checked_out": checked_out,
            "overflow": pool.overflow(),
            "checked_in": pool.checkedin(),
            "utilization": round(utilization, 2),
            "max_overflow": pool._max_overflow,
        }
    }

# mcpgateway/config.py

# Pool health monitoring
pool_metrics_enabled: bool = Field(
    default=True,
    description="Enable Prometheus metrics for connection pool"
)
pool_metrics_interval: float = Field(
    default=5.0,
    description="Interval in seconds to update pool metrics"
)

# Readiness probe configuration
readiness_pool_threshold: float = Field(
    default=0.85,
    ge=0.0,
    le=1.0,
    description="Pool utilization threshold for readiness (0.85 = fail at 85%)"
)
readiness_check_db_connection: bool = Field(
    default=False,
    description="Perform actual DB query in readiness check (adds latency)"
)