"""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 _capacity_by_hypervisor(self, attr: str) -> int:
        """Sum `attr` over alive nodes, deduplicated per hypervisor.

        On hypervisors that pass ALL physical GPUs through to every LXC node,
        each node's nvidia-smi reports the SAME physical GPUs. Summing across
        nodes would multiply one GPU set by the node count (e.g. 11 nodes x 8
        GPUs => ~4.3 TB phantom capacity). We therefore count each hypervisor
        ONCE (max over its nodes). Nodes whose hypervisor is unknown can't be
        proven to share GPUs, so each is counted on its own (keyed by node_id).
        """
        best: dict[str, int] = {}
        for n in self.nodes.values():
            if not n.alive:
                continue
            val = getattr(n, attr, 0) or 0
            hv = (n.hypervisor or "").strip().lower()
            key = hv if hv and hv not in ("unknown", "?") else f"__node__{n.node_id}"
            if val > best.get(key, 0):
                best[key] = val
        return sum(best.values())

    def total_gpu_capacity_mb(self) -> int:
        """Total physical-GPU VRAM (hardware capacity), deduplicated per hypervisor."""
        return self._capacity_by_hypervisor("vram_total_mb")

    def total_gpu_free_mb(self) -> int:
        """Free physical-GPU VRAM (placement headroom), deduplicated per hypervisor."""
        return self._capacity_by_hypervisor("vram_free_mb")