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