"""Health, readiness, and metrics endpoints."""

from fastapi import APIRouter
from fastapi.responses import PlainTextResponse

from prometheus_client import (
    CollectorRegistry, Gauge, Counter, generate_latest,
)

from radix_edge.db import get_db

router = APIRouter()

# --- Prometheus metrics ---

registry = CollectorRegistry()

# GPU metrics
gpu_utilization = Gauge("radix_gpu_utilization_pct", "GPU utilization %", ["gpu_index"], registry=registry)
gpu_memory_used = Gauge("radix_gpu_memory_used_mb", "GPU memory used (MB)", ["gpu_index"], registry=registry)
gpu_temperature = Gauge("radix_gpu_temperature_celsius", "GPU temperature (C)", ["gpu_index"], registry=registry)
gpu_power_draw = Gauge("radix_gpu_power_draw_watts", "GPU power draw (W)", ["gpu_index"], registry=registry)

# Job metrics
jobs_total = Gauge("radix_jobs_total", "Total jobs by status", ["status"], registry=registry)
job_queue_depth = Gauge("radix_job_queue_depth", "Number of queued jobs", registry=registry)

# Scheduler metrics
scheduler_decisions = Counter("radix_scheduler_decisions_total", "Scheduler decisions made", registry=registry)
scheduler_exploration_ratio = Gauge("radix_scheduler_exploration_ratio", "Current exploration ratio", registry=registry)

# FEP metrics
fep_free_energy = Gauge("radix_fep_free_energy", "Current variational free energy", registry=registry)
fep_latent_state = Gauge("radix_fep_latent_state_prob", "Latent state probability", ["state"], registry=registry)

# Brain weights
brain_weight = Gauge("radix_brain_weight", "Brain scheduling weight", ["dimension"], registry=registry)


def _update_metrics():
    """Update all Prometheus metrics from current database state."""
    try:
        with get_db() as db:
            # GPU metrics
            gpu_rows = db.execute("SELECT * FROM gpus").fetchall()
            for g in gpu_rows:
                idx = str(g["gpu_index"])
                gpu_utilization.labels(gpu_index=idx).set(g["utilization_pct"])
                gpu_memory_used.labels(gpu_index=idx).set(g["memory_used_mb"])
                gpu_temperature.labels(gpu_index=idx).set(g["temperature_c"])
                gpu_power_draw.labels(gpu_index=idx).set(g["power_draw_w"])

            # Job metrics
            status_counts = db.execute(
                "SELECT status, COUNT(*) as c FROM jobs GROUP BY status"
            ).fetchall()
            for r in status_counts:
                jobs_total.labels(status=r["status"]).set(r["c"])

            queued = db.execute("SELECT COUNT(*) as c FROM jobs WHERE status = 'queued'").fetchone()
            job_queue_depth.set(queued["c"])

            # Brain weights
            brain = db.execute("SELECT * FROM brain_params WHERE key = 'current'").fetchone()
            if brain:
                brain_weight.labels(dimension="queue_depth").set(brain["w_queue_depth"])
                brain_weight.labels(dimension="gpu_saturation").set(brain["w_gpu_saturation"])
                brain_weight.labels(dimension="node_pressure").set(brain["w_node_pressure"])

            # FEP latest decision
            fep = db.execute(
                "SELECT * FROM fep_decisions ORDER BY decided_at DESC LIMIT 1"
            ).fetchone()
            if fep:
                fep_free_energy.set(fep["free_energy"])
                import json
                probs = json.loads(fep["state_probs"]) if fep["state_probs"] else {}
                for state, prob in probs.items():
                    fep_latent_state.labels(state=state).set(prob)

    except Exception:
        pass  # Metrics endpoint should never fail

    # Scheduler model metrics
    try:
        from radix_edge.scheduler.scheduler import get_model
        model = get_model()
        metrics = model.get_metrics()
        # Counter can only go up, so use a gauge for current total
        scheduler_exploration_ratio.set(metrics.get("exploration_ratio", 0))
    except Exception:
        pass


# --- Endpoints ---

@router.get("/healthz")
async def health_check():
    """Liveness probe."""
    return {"status": "healthy"}


@router.get("/readyz")
async def readiness_check():
    """Readiness probe — verifies database is accessible."""
    try:
        with get_db() as db:
            db.execute("SELECT 1")
        return {"status": "ready"}
    except Exception as e:
        return {"status": "not_ready", "error": str(e)}, 503


@router.get("/metrics")
async def metrics():
    """Prometheus metrics endpoint."""
    _update_metrics()
    return PlainTextResponse(
        generate_latest(registry).decode("utf-8"),
        media_type="text/plain; version=0.0.4; charset=utf-8",
    )
