orochi-discord/orochi/bot.py

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

import disnake
from disnake import CategoryChannel, PermissionOverwrite, TextChannel
from disnake.ext import commands
import logging

from orochi.config import Config

bot = commands.Bot(command_prefix='!')
GAME: "Game"


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


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.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
    vote: Vote
    timestamp: datetime

    @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.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:
    rounds: list[Round] = field(default_factory=list)
    players: dict[str, Player] = field(default_factory=dict)

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

    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:
                return pickle.load(f)
        except FileNotFoundError:
            return None


@bot.event
async def on_ready():
    global GAME

    config: Config = bot.config
    logger = bot.logger

    if config.guild is None:
        config.save()
        logger.error("The guild ID is missing")
        exit(1)

    guild = await bot.fetch_guild(config.guild)

    if not guild:
        logger.error("Unknown guild.")
        exit(1)

    if config.vote_category is None:
        category = await guild.create_category("Votes")
        config.vote_category = category.id
        config.save()

    if config.secret_category is None:
        category = await guild.create_category("Conversation⋅s secrète⋅s")
        config.secret_category = category.id
        config.save()

    vote_category: CategoryChannel = await guild.fetch_channel(config.vote_category)
    if vote_category is None:
        config.vote_category = None
        return await on_ready()

    secret_category: CategoryChannel = await guild.fetch_channel(config.secret_category)
    if secret_category is None:
        config.secret_category = None
        return await on_ready()

    await vote_category.set_permissions(
        guild.default_role, overwrite=PermissionOverwrite(read_message_history=False, read_messages=False)
    )

    await secret_category.set_permissions(
        guild.default_role, overwrite=PermissionOverwrite(read_message_history=False, read_messages=False)
    )

    for i, player in enumerate(Config.PLAYERS):
        player_id = player.lower()
        if player_id not in config.vote_channels:
            channel: TextChannel = await vote_category.create_text_channel(player_id)
            config.vote_channels[player_id] = channel.id
            config.save()

        channel: TextChannel = await guild.fetch_channel(config.vote_channels[player_id])
        if channel is None:
            del config.vote_channels[player_id]
            return await on_ready()

        await channel.edit(name=player_id, category=vote_category, position=i)
        await channel.set_permissions(
            guild.default_role, overwrite=PermissionOverwrite(read_message_history=False, read_messages=False)
        )

        if player_id not in config.player_roles:
            role = await guild.create_role(name=player)
            config.player_roles[player_id] = role.id
            config.save()

        guild = await bot.fetch_guild(guild.id)  # update roles
        role = guild.get_role(config.player_roles[player_id])
        if role is None:
            del config.player_roles[player_id]
            config.save()
            return await on_ready()

        await channel.set_permissions(
            role, overwrite=PermissionOverwrite(read_message_history=True, read_messages=True)
        )

    GAME = Game.load('game.save')
    if not GAME:
        GAME = Game()
        for player in config.PLAYERS:
            GAME.register_player(player, config.vote_channels[player.lower()])
        GAME.save('game.save')

    # Update private channel id if necessary
    for player in list(GAME.players.values()):
        if player.private_channel_id != config.vote_channels[player.name.lower()]:
            GAME.register_player(player.name, config.vote_channels[player.name.lower()])
            GAME.save('game.save')

    if not config.telepathy_channel:
        channel: TextChannel = await secret_category.create_text_channel("bigbrain")
        config.telepathy_channel = channel.id
        config.save()

    telepathy_channel: TextChannel = await guild.fetch_channel(config.telepathy_channel)
    if not telepathy_channel:
        config.telepathy_channel = None
        return await on_ready()

    await telepathy_channel.edit(name="bigbrain", category=secret_category, position=0,
                                 topic="Échanges télépathiques")
    await telepathy_channel.set_permissions(
        guild.default_role, overwrite=PermissionOverwrite(read_message_history=False, read_messages=False)
    )
    delphine = guild.get_role(config.player_roles['delphine'])
    philia = guild.get_role(config.player_roles['philia'])
    await telepathy_channel.set_permissions(
        delphine, overwrite=PermissionOverwrite(read_message_history=True, read_messages=True)
    )
    await telepathy_channel.set_permissions(
        philia, overwrite=PermissionOverwrite(read_message_history=True, read_messages=True)
    )

    if not config.brother_channel:
        channel: TextChannel = await secret_category.create_text_channel("doliprane")
        config.brother_channel = channel.id
        config.save()

    brother_channel: TextChannel = await guild.fetch_channel(config.brother_channel)
    if not brother_channel:
        config.brother_channel = None
        return await on_ready()

    await brother_channel.edit(name="doliprane", category=secret_category, position=1,
                               topic="Des voix dans la tête ...")
    await brother_channel.set_permissions(
        guild.default_role, overwrite=PermissionOverwrite(read_message_history=False, read_messages=False)
    )
    await brother_channel.set_permissions(
        philia, overwrite=PermissionOverwrite(read_message_history=True, read_messages=True)
    )

    brother_channel_webhook = None
    if config.brother_channel_webhook is not None:
        try:
            brother_channel_webhook = await bot.fetch_webhook(config.brother_channel_webhook)
        except disnake.HTTPException | disnake.NotFound | disnake.Forbidden:
            pass
    if brother_channel_webhook is None:
        brother_channel_webhook = await brother_channel.create_webhook(name="???")
        config.brother_channel_webhook = brother_channel_webhook.id
        config.save()

    if not config.backdoor_channel:
        channel: TextChannel = await secret_category.create_text_channel("backdoor")
        config.backdoor_channel = channel.id
        config.save()

    backdoor_channel: TextChannel = await guild.fetch_channel(config.backdoor_channel)
    if not backdoor_channel:
        config.backdoor_channel = None
        return await on_ready()

    await backdoor_channel.edit(name="backdoor", category=secret_category, position=2,
                                topic="Panel d'administrati0n du jeu")
    await backdoor_channel.set_permissions(
        guild.default_role, overwrite=PermissionOverwrite(read_message_history=False, read_messages=False)
    )
    dan = guild.get_role(config.player_roles['dan'])
    await backdoor_channel.set_permissions(
        dan, overwrite=PermissionOverwrite(read_message_history=True, read_messages=True)
    )

    config.save()


@bot.command(help="Envoyer un message en tant qu'Orochi.")
@commands.has_permissions(administrator=True)
async def send(ctx: commands.Context, *, message: str):
    await ctx.message.delete()
    await ctx.send(message)


@bot.command(help="Envoyer un message à Philia par la pensée en tant que Brother.")
@commands.has_permissions(administrator=True)
async def brother(ctx: commands.Context, *, message: str):
    webhook = await bot.fetch_webhook(bot.config.brother_channel_webhook)
    await webhook.send(message)
    await ctx.message.reply("Message envoyé.")


@bot.command()
async def vote(ctx: commands.Context):
    view = Confirm()
    await ctx.message.reply("plop", view=view)
    await view.wait()


# Define a simple View that gives us a confirmation menu
class Confirm(disnake.ui.View):
    @disnake.ui.button(label="S'allier", style=disnake.ButtonStyle.green)
    async def confirm(self, button: disnake.ui.Button, interaction: disnake.MessageInteraction):
        self.clear_items()
        await interaction.response.edit_message(content="Vous vous êtes allié.", view=self)
        self.stop()

    @disnake.ui.button(label="Trahir", style=disnake.ButtonStyle.red)
    async def cancel(self, button: disnake.ui.Button, interaction: disnake.MessageInteraction):
        self.clear_items()
        await interaction.response.edit_message(content="Vous avez trahi.", view=self)
        self.stop()


def run():
    config = Config.load()

    logger = logging.getLogger('discord')
    logger.setLevel(logging.DEBUG)
    handler = logging.FileHandler(filename='../discord.log', encoding='utf-8', mode='w')
    handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s:%(name)s: %(message)s'))
    logger.addHandler(handler)

    bot.config = config
    bot.logger = logger

    bot.run(config.discord_token)