"""Heartbeat background loop for upstream control plane."""

from __future__ import annotations

import asyncio
import logging
import os

from radix_edge.config import get_config
from radix_edge.db import get_db
from radix_edge.upstream.client import UpstreamClient

logger = logging.getLogger("radix_edge.upstream")


async def heartbeat_loop(client: UpstreamClient) -> None:
    """Send periodic heartbeats to control plane."""
    config = get_config()
    logger.info("Upstream heartbeat started (interval=%.0fs)", config.heartbeat_interval)

    while True:
        try:
            capabilities = _build_capabilities()
            await client.send_heartbeat(status="active", capabilities=capabilities)
            logger.debug("Heartbeat sent")
        except Exception:
            logger.exception("Heartbeat error")

        await asyncio.sleep(config.heartbeat_interval)


def _build_capabilities() -> dict:
    """Build capabilities dict from current GPU and job state."""
    capabilities: dict = {
        "cpu_count": os.cpu_count() or 1,
    }

    try:
        with get_db() as db:
            gpu_rows = db.execute(
                "SELECT gpu_index, name, memory_total_mb FROM gpus ORDER BY gpu_index"
            ).fetchall()
            gpus = []
            for r in gpu_rows:
                gpus.append({"name": r["name"], "memory": f"{r['memory_total_mb']} MiB"})
            capabilities["gpus"] = gpus
            capabilities["gpu_count"] = len(gpus)

            running = db.execute(
                "SELECT COUNT(*) as c FROM jobs WHERE status IN ('running', 'scheduled')"
            ).fetchone()
            capabilities["running_jobs"] = running["c"]

            queued = db.execute(
                "SELECT COUNT(*) as c FROM jobs WHERE status = 'queued'"
            ).fetchone()
            capabilities["queued_jobs"] = queued["c"]
    except Exception:
        capabilities["gpus"] = []
        capabilities["gpu_count"] = 0

    return capabilities
