Persistent server-side group assignments

services/server_state.py

@@ -0,0 +1,139 @@
+# services/server_state.py
+from typing import Dict, List, Optional
+from pydantic import BaseModel
+import asyncio
+
+class Member(BaseModel):
+    id: int
+    name: str
+    level: int
+    estimate: str
+    type: str                 # "friendly" or "enemy"
+    group: Optional[str] = None   # "group1", "group2", ... or None
+    hits: int = 0
+
+class ServerState:
+    def __init__(self, group_count: int = 5):
+        # member maps
+        self.friendly: Dict[int, Member] = {}
+        self.enemy: Dict[int, Member] = {}
+        # initialize groups 1..group_count as strings "1","2",...
+        self.group_count = group_count
+        self.groups: Dict[str, Dict[str, List[int]]] = {}
+        for i in range(1, group_count + 1):
+            key = str(i)
+            self.groups[key] = {"friendly": [], "enemy": []}
+
+        # bot running flag
+        self.bot_running: bool = False
+
+        # concurrency lock for async safety
+        self.lock = asyncio.Lock()
+
+    # Assignment helpers (all use the lock when called from async endpoints)
+    async def assign_member(self, member_id: int, kind: str, group_key: str):
+        """
+        Assign a member (by id) of kind ('friendly'|'enemy') to group_key ("1".."N").
+        This removes it from any previous group of the same kind.
+        """
+        async with self.lock:
+            if kind not in ("friendly", "enemy"):
+                raise ValueError("invalid kind")
+
+            if group_key not in self.groups:
+                raise ValueError("invalid group_key")
+
+            # remove from any existing group
+            for gk, buckets in self.groups.items():
+                if member_id in buckets[kind]:
+                    buckets[kind].remove(member_id)
+
+            # add to new group if not present
+            if member_id not in self.groups[group_key][kind]:
+                self.groups[group_key][kind].append(member_id)
+
+            # update member.group
+            coll = self.friendly if kind == "friendly" else self.enemy
+            if member_id in coll:
+                coll[member_id].group = group_key
+
+    async def remove_member_assignment(self, member_id: int):
+        """Remove member from any group (both friendly and enemy)."""
+        async with self.lock:
+            for gk, buckets in self.groups.items():
+                if member_id in buckets["friendly"]:
+                    buckets["friendly"].remove(member_id)
+                if member_id in buckets["enemy"]:
+                    buckets["enemy"].remove(member_id)
+
+            # clear group attr if member in maps
+            if member_id in self.friendly:
+                self.friendly[member_id].group = None
+            if member_id in self.enemy:
+                self.enemy[member_id].group = None
+
+    async def clear_all_assignments(self):
+        """Clear all group lists and member.group fields and save to disk."""
+        async with self.lock:
+            # clear in-memory groups
+            for gk in self.groups:
+                self.groups[gk]["friendly"].clear()
+                self.groups[gk]["enemy"].clear()
+            
+            # clear group for known members
+            for m in self.friendly.values():
+                m.group = None
+            for m in self.enemy.values():
+                m.group = None
+
+            # save to disk so clients get empty state
+            self.save_assignments()
+
+    async def get_assignments_snapshot(self) -> Dict[str, Dict[str, List[int]]]:
+        """Return a copy snapshot of the groups dictionary."""
+        async with self.lock:
+            # shallow copy of structure
+            snap = {}
+            for gk, buckets in self.groups.items():
+                snap[gk] = {
+                    "friendly": list(buckets["friendly"]),
+                    "enemy": list(buckets["enemy"])
+                }
+            return snap
+
+    # member add/update helpers
+    async def upsert_member(self, member_data: dict, kind: str):
+        """
+        Insert or update member. member_data expected to have id,name,level,estimate.
+        kind must be 'friendly' or 'enemy'.
+        """
+        async with self.lock:
+            if kind not in ("friendly", "enemy"):
+                raise ValueError("invalid kind")
+
+            coll = self.friendly if kind == "friendly" else self.enemy
+            mid = int(member_data["id"])
+            existing_group = coll.get(mid).group if (mid in coll) else None
+            member = Member(
+                id=mid,
+                name=member_data.get("name", "Unknown"),
+                level=int(member_data.get("level", 0)),
+                estimate=str(member_data.get("estimate", "?")),
+                type=kind,
+                group=existing_group,
+                hits=int(member_data.get("hits", 0)) if "hits" in member_data else 0
+            )
+            coll[mid] = member
+
+    async def remove_missing_members(self, received_ids: List[int], kind: str):
+        """Remove any existing members not present in received_ids (useful after a populate)."""
+        async with self.lock:
+            coll = self.friendly if kind == "friendly" else self.enemy
+            to_remove = [mid for mid in coll.keys() if mid not in set(received_ids)]
+            for mid in to_remove:
+                # remove from groups too
+                await self.remove_member_assignment(mid)
+                del coll[mid]
+
+# single global state (5 groups)
+STATE = ServerState(group_count=5)