node_router.py

"""
node_router.py — Server-side. Registers LLM nodes, assigns monsters,
executes incoming commands. Drop into server.py as a member of start_server().
"""

import json
import math
import time


class NodeRouter:
    """
    Receives 'register' packets from LLM nodes.
    Assigns monster accounts to nodes.
    Executes 'command' packets from nodes.
    Despawns accounts when a node goes silent.
    """

    NODE_TIMEOUT = 15.0  # Seconds before a node is considered dead

    def __init__(self, sock, clients, active_monsters, heightmap):
        self.sock            = sock
        self.clients         = clients
        self.active_monsters = active_monsters  # monster_manager.active_monsters ref
        self.heightmap       = heightmap

        # node_id → {addr, port, capacity, last_seen, assigned: set[monster_id]}
        self.nodes: dict[str, dict] = {}

        # monster_id → node_id  (reverse lookup for validation)
        self.monster_node_map: dict[str, str] = {}

    # ── Called from process_inbound ───────────────────────────────────────────
    def handle_packet(self, packet: dict, addr: tuple, current_time: float):
        ptype = packet.get("type")

        if ptype == "register":
            self._handle_register(packet, addr, current_time)

        elif ptype == "command":
            self._handle_command(packet, current_time)

    def _handle_register(self, packet: dict, addr: tuple, current_time: float):
        nid = packet["node_id"]
        if nid not in self.nodes:
            print(f"[Router] New LLM node: {nid} @ {addr}")
            self.nodes[nid] = {
                "addr":      addr,
                "port":      packet.get("port", addr[1]),
                "last_seen": current_time,
                "assigned":  set(),
            }
        else:
            self.nodes[nid]["last_seen"] = current_time

    def _handle_command(self, packet: dict, current_time: float):
        mid     = packet.get("monster_id")
        command = packet.get("command", "").upper()
        thought = packet.get("thought", "")
        nid     = packet.get("node_id")

        # Validate ownership
        if self.monster_node_map.get(mid) != nid:
            return  # Node doesn't own this monster — reject

        monster = self.active_monsters.get(mid)
        if not monster:
            return

        if command not in {"ATTACK", "FLEE", "STALK", "WANDER", "IGNORE"}:
            return

        monster.state        = command
        monster.thought_text = thought
        monster.is_thinking  = False
        print(f"[Router] {mid} → {command} | \"{thought}\"")

    # ── Called from process_outbound ──────────────────────────────────────────
    def tick(self, current_time: float):
        self._cull_dead_nodes(current_time)
        self._assign_unrouted_monsters()
        self._broadcast_player_snapshot()

    def _cull_dead_nodes(self, current_time: float):
        dead = [
            nid for nid, nd in self.nodes.items()
            if current_time - nd["last_seen"] > self.NODE_TIMEOUT
        ]
        for nid in dead:
            assigned = self.nodes[nid]["assigned"]
            print(f"[Router] Node {nid} timed out — despawning "
                  f"{len(assigned)} monster accounts.")
            for mid in assigned:
                self.active_monsters.pop(mid, None)
                self.monster_node_map.pop(mid, None)
            del self.nodes[nid]

    def _assign_unrouted_monsters(self):
        """Pushes monsters not yet on any node to the least-loaded available node."""
        if not self.nodes:
            return

        unrouted = [
            mid for mid in self.active_monsters
            if mid not in self.monster_node_map
        ]
        if not unrouted:
            return

        # Pick least-loaded node
        target_nid = min(
            self.nodes,
            key=lambda nid: len(self.nodes[nid]["assigned"])
        )
        node = self.nodes[target_nid]

        monsters_out = []
        for mid in unrouted:
            m = self.active_monsters[mid]
            self.monster_node_map[mid] = target_nid
            node["assigned"].add(mid)
            monsters_out.append({
                "id":          mid,
                "personality": m.personality_type,
                "x": m.x, "y": m.y, "z": m.z,
            })

        self._send_to_node(target_nid, {
            "type":     "assign",
            "monsters": monsters_out,
        })
        print(f"[Router] Assigned {len(monsters_out)} monsters → node {target_nid}")

    def _broadcast_player_snapshot(self):
        """Sends current player positions to all nodes so they can self-manage evals."""
        if not self.nodes or not self.clients:
            return

        snapshot = {
            "type": "player_snapshot",
            "players": [
                {"id": cid, "x": cd["x"], "y": cd["y"], "z": cd["z"],
                 "hp": cd.get("hp", 100), "item": cd.get("item", 0)}
                for cid, cd in self.clients.items() if "x" in cd
            ]
        }
        for nid in self.nodes:
            self._send_to_node(nid, snapshot)

    def _send_to_node(self, nid: str, packet: dict):
        node = self.nodes.get(nid)
        if not node:
            return
        try:
            self.sock.sendto(
                json.dumps(packet).encode(),
                (node["addr"][0], node["port"])
            )
        except Exception as e:
            print(f"[Router] Send to {nid} failed: {e}")