"""Background loops for mesh networking."""

from __future__ import annotations

import asyncio
import logging
import os
from typing import TYPE_CHECKING

from radix_edge.db import get_db

if TYPE_CHECKING:
    from radix_edge.config import EdgeConfig
    from radix_edge.mesh.coordinator import MeshCoordinator
    from radix_edge.mesh.discovery import MeshDiscovery, ManualPeerList
    from radix_edge.mesh.peer_client import PeerClient
    from radix_edge.mesh.peer_registry import PeerRegistry

logger = logging.getLogger("radix_edge.mesh.sync")


def _build_capabilities() -> dict:
    """Build capabilities dict from current GPU and job state.

    Reuses the same pattern as upstream/heartbeat.py.
    """
    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


async def mesh_heartbeat_loop(
    config: "EdgeConfig",
    registry: "PeerRegistry",
    peer_clients: dict[str, "PeerClient"],
    node_id: str,
    coordinator: "MeshCoordinator",
) -> None:
    """Send heartbeat to coordinator (or all peers if no coordinator yet)."""
    logger.info("Mesh heartbeat loop started (interval=%.1fs)", config.mesh_heartbeat_interval)

    while True:
        try:
            capabilities = _build_capabilities()
            capabilities["port"] = config.port

            if coordinator.coordinator_id and coordinator.coordinator_id != node_id:
                # Send to coordinator only
                cid = coordinator.coordinator_id
                if cid in peer_clients:
                    try:
                        await peer_clients[cid].send_heartbeat(node_id, capabilities)
                        logger.debug("Heartbeat sent to coordinator %s", cid)
                    except Exception:
                        logger.warning("Failed to send heartbeat to coordinator %s", cid)
            else:
                # No coordinator or we are coordinator: send to all peers
                for nid, client in peer_clients.items():
                    try:
                        await client.send_heartbeat(node_id, capabilities)
                    except Exception:
                        logger.debug("Failed to send heartbeat to peer %s", nid)

        except Exception:
            logger.exception("Mesh heartbeat error")

        await asyncio.sleep(config.mesh_heartbeat_interval)


async def mesh_discovery_loop(
    config: "EdgeConfig",
    discovery: "MeshDiscovery | ManualPeerList | None",
    registry: "PeerRegistry",
    peer_clients: dict[str, "PeerClient"],
    mesh_key: str,
) -> None:
    """Periodically scan for new peers, remove stale ones."""
    from radix_edge.mesh.peer_client import PeerClient

    logger.info("Mesh discovery loop started (interval=%.1fs)", config.mesh_discovery_interval)

    while True:
        try:
            # Update from discovery source
            if hasattr(discovery, "discovered_peers"):
                # mDNS discovery
                for nid, info in discovery.discovered_peers.items():
                    if nid not in peer_clients:
                        peer_clients[nid] = PeerClient(info["host"], info["port"], mesh_key)
                        logger.info("New peer client created for %s at %s:%d", nid, info["host"], info["port"])
            elif hasattr(discovery, "peers"):
                # Manual peer list — peers don't have node_ids until they heartbeat
                for i, peer_info in enumerate(discovery.peers):
                    placeholder_id = f"manual-{peer_info['host']}:{peer_info['port']}"
                    if placeholder_id not in peer_clients:
                        peer_clients[placeholder_id] = PeerClient(peer_info["host"], peer_info["port"], mesh_key)

            # Remove stale peers
            for stale in registry.get_stale_peers():
                registry.remove_peer(stale.node_id)
                if stale.node_id in peer_clients:
                    client = peer_clients.pop(stale.node_id)
                    await client.close()
                logger.info("Removed stale peer: %s", stale.node_id)

        except Exception:
            logger.exception("Mesh discovery error")

        await asyncio.sleep(config.mesh_discovery_interval)


async def election_monitor_loop(
    config: "EdgeConfig",
    coordinator: "MeshCoordinator",
    registry: "PeerRegistry",
    peer_clients: dict[str, "PeerClient"],
) -> None:
    """Watch coordinator liveness. Trigger re-election if stale."""
    logger.info("Election monitor started (check interval=%.1fs)", config.mesh_stale_threshold)

    while True:
        try:
            cid = coordinator.coordinator_id

            if cid and cid != coordinator.node_id:
                # Check if coordinator is still alive
                peer = registry.get_peer(cid)
                if peer is None or peer in registry.get_stale_peers():
                    logger.warning("Coordinator %s is stale, triggering re-election", cid)
                    new_coordinator = coordinator.elect()

                    # Announce to all peers
                    for nid, client in peer_clients.items():
                        try:
                            await client.announce_election(new_coordinator)
                        except Exception:
                            logger.debug("Failed to announce election to peer %s", nid)
            elif cid is None:
                # No coordinator yet, run election
                new_coordinator = coordinator.elect()
                for nid, client in peer_clients.items():
                    try:
                        await client.announce_election(new_coordinator)
                    except Exception:
                        logger.debug("Failed to announce election to peer %s", nid)

        except Exception:
            logger.exception("Election monitor error")

        await asyncio.sleep(config.mesh_stale_threshold)


async def coordinator_loop(
    config: "EdgeConfig",
    coordinator: "MeshCoordinator",
    peer_clients: dict[str, "PeerClient"],
) -> None:
    """Coordinator-only: run distributed scheduling at 2s intervals."""
    logger.info("Coordinator scheduling loop started")

    while True:
        try:
            if coordinator.is_coordinator:
                assigned = await coordinator.schedule_across_mesh(peer_clients)
                if assigned > 0:
                    logger.info("Coordinator assigned %d jobs across mesh", assigned)
        except Exception:
            logger.exception("Coordinator scheduling error")

        await asyncio.sleep(2.0)
