121 lines
4.5 KiB
Python
121 lines
4.5 KiB
Python
# 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 # "1", "2", ...
|
|
hits: int = 0
|
|
status: str = "Unknown" # Added status field for Torn API status
|
|
|
|
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
|
|
async def assign_member(self, member_id: int, kind: str, group_key: str):
|
|
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 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):
|
|
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 attribute
|
|
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):
|
|
async with self.lock:
|
|
for gk in self.groups:
|
|
self.groups[gk]["friendly"].clear()
|
|
self.groups[gk]["enemy"].clear()
|
|
for m in self.friendly.values():
|
|
m.group = None
|
|
for m in self.enemy.values():
|
|
m.group = None
|
|
|
|
async def get_assignments_snapshot(self) -> Dict[str, Dict[str, List[int]]]:
|
|
async with self.lock:
|
|
snap = {}
|
|
for gk, buckets in self.groups.items():
|
|
snap[gk] = {
|
|
"friendly": list(buckets["friendly"]),
|
|
"enemy": list(buckets["enemy"])
|
|
}
|
|
return snap
|
|
|
|
# member helpers
|
|
async def upsert_member(self, member_data: dict, kind: str):
|
|
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,
|
|
status=member_data.get("status", "Unknown") # Initialize status if provided
|
|
)
|
|
coll[mid] = member
|
|
|
|
async def remove_missing_members(self, received_ids: List[int], kind: str):
|
|
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:
|
|
await self.remove_member_assignment(mid)
|
|
del coll[mid]
|
|
|
|
# Single global state
|
|
STATE = ServerState(group_count=5)
|