support/manage_game.py

import random

from support.game_settings import CODENAMES_WORDS, NAMES, TEAM_MODEL_PRESETS, ALL_MODELS
from support.load_models import create_model
from support.settings import AVATAR_PATHS
from typing import List, Optional


class Player:
    def __init__(self, name: str, model_name: str, model, avatar: str):
        self.name = name
        self.model_name = model_name
        self.model = model
        self.avatar = avatar
        self.team: Optional[str] = None
        self.role: Optional[str] = None


class Game:
    def __init__(
        self,
        selected_team: str = None,
        red_team: str = "google",
        blue_team: str = "openai",
        api_keys: dict = None,
    ):
        """
        Initialize a new game.
        Either provide `selected_team` (legacy mode) or `red_team` and `blue_team`.

        - red_team / blue_team: choose from ["openai", "google", "anthropic", "opensource", "random"]
        - selected_team: legacy single-team mode (both teams use same preset)
        """
        # Backward compatibility
        if selected_team and not (red_team or blue_team):
            self.red_team_choice = selected_team.lower()
            self.blue_team_choice = selected_team.lower()
        else:
            self.red_team_choice = red_team.lower()
            self.blue_team_choice = blue_team.lower()

        self.players: List[Player] = []
        self.board: dict = {}
        self.api_keys = api_keys if api_keys else {}
        self.setup_game()

    def setup_game(self):
        """Initialize a new game with players and board"""
        # Prepare shared pools for names and avatars
        available_names = NAMES.copy()
        available_avatars = AVATAR_PATHS.copy()
        random.shuffle(available_names)
        random.shuffle(available_avatars)

        self.players = self.create_players(available_names, available_avatars)
        self.board = self.generate_board()

    def _generate_team_players(self, team_color, team_choice, available_names, available_avatars):
        players = []
        roles = ["boss", "captain", "player", "player"]

        if team_choice in TEAM_MODEL_PRESETS:
            setup = TEAM_MODEL_PRESETS[team_choice]
            model_names = [setup["boss"], setup["captain"]] + setup["players"]
        else:
            ROLE_POOLS = {
                "boss": [preset["boss"] for preset in TEAM_MODEL_PRESETS.values()],
                "captain": [preset["captain"] for preset in TEAM_MODEL_PRESETS.values()],
                "player": [
                    model
                    for preset in TEAM_MODEL_PRESETS.values()
                    for model in preset["players"]
                ]
            }
            model_names = [random.choice(ROLE_POOLS[role]) for role in roles]

        for i in range(4):
            name = available_names.pop() if available_names else f"{team_color.title()}_{i}"
            avatar = available_avatars.pop() if available_avatars else "assets/avatars/default.png"
            model_name = model_names[i]
            model = create_model(model_name, self.api_keys)
            p = Player(name, model_name, model, avatar)
            p.team = team_color
            p.role = roles[i]
            players.append(p)
        return players

    def create_players(self, available_names, available_avatars) -> List[Player]:
        """Create 8 players (4 per team) according to selected team presets"""
        red_players = self._generate_team_players("red", self.red_team_choice, available_names, available_avatars)
        blue_players = self._generate_team_players("blue", self.blue_team_choice, available_names, available_avatars)
        return red_players + blue_players

    def generate_board(self) -> dict:
        """
        Generate a new game board with 25 random words.
        Distribution: 9 starting team, 8 other team, 7 neutral, 1 killer
        """
        # Select 25 random words
        selected_words = random.sample(CODENAMES_WORDS, 25)

        # Randomly choose starting team
        starting_team = random.choice(['red', 'blue'])

        # Assign word counts based on starting team
        if starting_team == 'red':
            red_count, blue_count = 9, 8
        else:
            red_count, blue_count = 8, 9

        neutral_count = 7
        # Create a list of (word, color) tuples
        word_color_pairs = []
        idx = 0

        # Add red words
        for i in range(red_count):
            word_color_pairs.append((selected_words[idx], 'red'))
            idx += 1

        # Add blue words
        for i in range(blue_count):
            word_color_pairs.append((selected_words[idx], 'blue'))
            idx += 1

        # Add neutral words
        for i in range(neutral_count):
            word_color_pairs.append((selected_words[idx], 'neutral'))
            idx += 1

        # Add killer word
        word_color_pairs.append((selected_words[idx], 'killer'))

        # Shuffle to randomize positions on board
        random.shuffle(word_color_pairs)

        # Store in board dictionary
        board = {
            'word_color_pairs': word_color_pairs,
            'starting_team': starting_team,
            'red': [w for w, c in word_color_pairs if c == 'red'],
            'blue': [w for w, c in word_color_pairs if c == 'blue'],
            'neutral': [w for w, c in word_color_pairs if c == 'neutral'],
            'killer': [w for w, c in word_color_pairs if c == 'killer'][0]
        }

        return board