"""In-memory registry for nodes and modules."""
from __future__ import annotations
import time
import logging

from fabric.shared.schemas import (
    Heartbeat, NodeInfo, ModuleStatus, ModuleState,
    PlaceRequest, PlaceResponse, PriorityClass,
)

log = logging.getLogger("fabric.registry")

STALE_NODE_TIMEOUT_S = 60


class Registry:
    def __init__(self):
        self.nodes: dict[str, NodeInfo] = {}
        self._version = 0

    @property
    def version(self) -> int:
        return self._version

    def handle_heartbeat(self, hb: Heartbeat) -> None:
        node = self.nodes.get(hb.node_id)
        if node is None:
            log.info("New node registered: %s (%s)", hb.node_id, hb.node_ip)
        self.nodes[hb.node_id] = NodeInfo(
            node_id=hb.node_id,
            node_ip=hb.node_ip,
            agent_port=hb.agent_port,
            hypervisor=hb.hypervisor,
            gpu_status=hb.gpu_status,
            modules=hb.modules,
            last_heartbeat=time.time(),
            alive=True,
            gpus=hb.gpus,
            vram_total_mb=hb.vram_total_mb,
            vram_free_mb=hb.vram_free_mb,
            driver_version=hb.driver_version,
        )
        self._version += 1

    def prune_stale_nodes(self) -> list[str]:
        now = time.time()
        stale = []
        for nid, node in list(self.nodes.items()):
            if now - node.last_heartbeat > STALE_NODE_TIMEOUT_S:
                node.alive = False
                stale.append(nid)
        return stale

    def get_all_modules(self) -> list[ModuleStatus]:
        modules = []
        for node in self.nodes.values():
            if not node.alive:
                continue
            for m in node.modules:
                modules.append(m)
        return modules

    def get_alive_nodes(self) -> list[NodeInfo]:
        return [n for n in self.nodes.values() if n.alive]

    def find_module(self, model_name: str) -> list[tuple[NodeInfo, ModuleStatus]]:
        results = []
        for node in self.nodes.values():
            if not node.alive:
                continue
            for m in node.modules:
                if m.name == model_name or model_name in m.aliases:
                    results.append((node, m))
        return results

    def place(self, req: PlaceRequest) -> PlaceResponse | None:
        candidates = self.find_module(req.model)
        if not candidates:
            return None

        state_priority = {
            ModuleState.running: 0,
            ModuleState.loaded: 1,
            ModuleState.idle: 2,
            ModuleState.sleeping: 3,
            ModuleState.warming: 4,
            ModuleState.unloaded: 5,
            ModuleState.error: 6,
        }
        # Capacity-aware tie-break: among equally-good states, prefer the node
        # with the most free physical VRAM (more headroom => safer to (re)load).
        candidates.sort(key=lambda c: (state_priority.get(c[1].state, 99), -(c[0].vram_free_mb or 0)))

        node, module = candidates[0]
        return PlaceResponse(
            node_id=node.node_id,
            node_ip=node.node_ip,
            infer_url=module.infer_url,
            module_name=module.name,
            state=module.state,
        )

    def total_gpu_vram_mb(self) -> int:
        total = 0
        for node in self.nodes.values():
            if not node.alive:
                continue
            for g in node.gpu_status:
                total += g.vram_mb
        return total

    def total_gpu_capacity_mb(self) -> int:
        """Sum of physical-GPU total VRAM across alive nodes (hardware capacity)."""
        return sum(n.vram_total_mb for n in self.nodes.values() if n.alive)

    def total_gpu_free_mb(self) -> int:
        """Sum of physical-GPU free VRAM across alive nodes (placement headroom)."""
        return sum(n.vram_free_mb for n in self.nodes.values() if n.alive)