orochi-discord/orochi/models.py

import pickle
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import ClassVar


class Room(Enum):
    A = 'A'
    B = 'B'
    C = 'C'


class GameState(Enum):
    PREPARING = 0
    VOTING = 1
    RESULTS = 2


class Vote(Enum):
    ALLY = 'A'
    BETRAY = 'B'


@dataclass(frozen=True)
class Player:
    name: str
    private_channel_id: int = field(hash=False)

    @property
    def round_votes(self):
        for r in Game.INSTANCE.rounds:
            for room in r.rooms:
                for vote in room.votes:
                    if self in vote.players:
                        yield vote

    @property
    def score(self):
        s = 3

        for vote in self.round_votes:
            room = vote.room
            other_vote = room.vote1 if room.vote1 is not vote else room.vote2
            match vote.vote, other_vote.vote:
                case Vote.ALLY, Vote.ALLY:
                    s += 2
                case Vote.ALLY, Vote.BETRAY:
                    s -= 2
                case Vote.BETRAY, Vote.ALLY:
                    s += 3
                case Vote.BETRAY, Vote.BETRAY:
                    pass

        return s


@dataclass
class RoundVote:
    player1: Player
    player2: Player | None = None
    vote: Vote | None = None
    timestamp: datetime | None = None

    @property
    def players(self):
        return self.player1, self.player2

    @property
    def room(self):
        for r in Game.rounds:
            for room in r.rooms:
                if self in room.votes:
                    return room


@dataclass
class RoundRoom:
    room: Room
    vote1: RoundVote
    vote2: RoundVote

    @property
    def votes(self):
        return self.vote1, self.vote2

    @property
    def round(self):
        for r in Game.INSTANCE.rounds:
            if self in r.rooms:
                return r


@dataclass
class Round:
    round: int
    room_a: RoundRoom
    room_b: RoundRoom
    room_c: RoundRoom

    @property
    def rooms(self):
        return self.room_a, self.room_b, self.room_c


@dataclass
class Game:
    INSTANCE: ClassVar["Game"] = None

    state: GameState = GameState.PREPARING
    rounds: list[Round] = field(default_factory=list)
    players: dict[str, Player] = field(default_factory=dict)
    
    def __post_init__(self):
        Game.INSTANCE = self

    def register_player(self, name: str, vote_channel_id: int) -> Player:
        player = Player(name, vote_channel_id)
        self.players[name] = player
        return player

    def default_first_round(self) -> Round:
        return Round(
            round=1,
            room_a=RoundRoom(room=Room.A,
                             vote1=RoundVote(player1=self.players['Tora']),
                             vote2=RoundVote(player1=self.players['Kamui'],
                                             player2=self.players['Philia'])),
            room_b=RoundRoom(room=Room.A,
                             vote1=RoundVote(player1=self.players['Dan']),
                             vote2=RoundVote(player1=self.players['Ennea'],
                                             player2=self.players['Delphine'])),
            room_c=RoundRoom(room=Room.A,
                             vote1=RoundVote(player1=self.players['Hanabi']),
                             vote2=RoundVote(player1=self.players['Nona'],
                                             player2=self.players['Oji'])),
        )

    def save(self, filename: str) -> None:
        """
        Uses pickle to save the current state of the game.
        """
        with open(filename, 'wb') as f:
            pickle.dump(self, f)

    @classmethod
    def load(cls, filename: str) -> "Game | None":
        """
        Reload the game from a saved file.
        """
        try:
            with open(filename, 'rb') as f:
                game = pickle.load(f)
                Game.INSTANCE = game
                return game
        except FileNotFoundError:
            return None