# 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)