orochi-discord/orochi/models.py

import pickle
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import ClassVar, Iterable, Generator


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) -> Generator["RoundVote", None, None]:
        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 calculated_score(self) -> int:
        s = 3

        for vote in self.round_votes:
            if vote.room.round.round == len(Game.INSTANCE.rounds) and Game.INSTANCE.state != GameState.RESULTS:
                # Don't compute temporary scores
                break

            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

            if s <= 0:
                # Player died
                return s

        return s

    @property
    def score(self) -> int:
        if self in Game.INSTANCE.score_overrides:
            return Game.INSTANCE.score_overrides[self]
        return self.calculated_score

    def __str__(self):
        return self.name


@dataclass
class RoundVote:
    player1: Player | None = None
    player2: Player | None = None
    vote: Vote | None = None
    swapped: bool = field(default=False, init=False)

    @property
    def players(self) -> Iterable[Player]:
        if self.player2 is None:
            return self.player1,
        return self.player1, self.player2

    @property
    def room(self) -> "RoundRoom":
        for r in Game.INSTANCE.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 players(self) -> Generator[Player, None, None]:
        for vote in self.votes:
            yield from vote.players

    @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) -> tuple[RoundRoom, RoundRoom, RoundRoom]:
        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)
    score_overrides: dict[Player, int] = field(default_factory=dict, init=False)
    
    def __post_init__(self: "Game") -> None:
        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.B,
                             vote1=RoundVote(player1=self.players['Dan']),
                             vote2=RoundVote(player1=self.players['Ennea'],
                                             player2=self.players['Delphine'])),
            room_c=RoundRoom(room=Room.C,
                             vote1=RoundVote(player1=self.players['Hanabi']),
                             vote2=RoundVote(player1=self.players['Nona'],
                                             player2=self.players['Oji'])),
        )

    def save(self, filename: str | None = None) -> None:
        """
        Uses pickle to save the current state of the game.
        """
        if filename is None:
            filename = Path(__file__).parent.parent / 'game.save'
        with open(filename, 'wb') as f:
            pickle.dump(self, f)

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