botmafieux/cogs/intern.py

import discord, json, os, random, re, datetime
import collections
from discord.ext import commands, bridge, tasks
from setup.logger import LOGGER
from setup.config import Config

class InternCog(commands.Cog):
    def __init__(self, bot:commands.Bot) -> None:
        self.bot = bot
        self.in_voice_channel = False
        self.connected_guild_ids = set()

    def _get_guild_data(self, guild_id):
        return next((data for data in self.data if data["guild_id"] == guild_id), None)

    def _get_guild_voice_client(self, guild_id):
        return next(
            (
                voice_client
                for voice_client in self.bot.voice_clients
                if getattr(getattr(voice_client, "guild", None), "id", None) == guild_id
            ),
            None,
        )

    def _resolve_voice_action(self, before_channel, after_channel, guild_voice_client):
        if after_channel is None and before_channel is not None:
            return "disconnect"

        if before_channel is not None and after_channel is not None:
            if before_channel.id != after_channel.id:
                return "move"
            return None

        if after_channel is not None and before_channel is None:
            if guild_voice_client is not None:
                return "move"
            return "connect" if random.random() <= 0.25 else None

        return None

    def load_from_file(self):
        try:
            with open(Config.INTERN_MEMORY_PATH) as f:
                self.data = json.load(f)
            if len(self.data) == 0: raise FileNotFoundError
            for guild_data in self.data:
                guild = self.bot.get_guild(guild_data["guild_id"])
                if guild:
                    mentor = guild.get_member(guild_data["mentor_id"])
                    LOGGER.debug(f"Loaded mentor: {mentor.name} for guild {guild.name}")
        except FileNotFoundError:
            today = datetime.datetime.now()
            next_roll = datetime.datetime(year=today.year,month=today.month,day=today.day) + datetime.timedelta(days=7)
            self.data = [
                {
                    "mentor_id": random.choice(guild.members).id,
                    "guild_id": guild.id,
                    "next_roll": next_roll.isoformat()
                }
                for guild in self.bot.guilds
            ]

    def load_from_database(self):
        pass

    def save_to_file(self):
        with open(Config.INTERN_MEMORY_PATH, 'w') as f:
            json.dump(self.data,f)

    def save_to_database(self):
        pass

    @tasks.loop(hours=24)
    async def update(self):
        today = datetime.datetime.now()
        for guild_data in self.data:
            next_roll = datetime.datetime.fromisoformat(guild_data["next_roll"])
            if today > next_roll:
                guild = self.bot.get_guild(guild_data["guild_id"])
                if guild:
                    members = [
                        member
                        for member in guild.get_role(1516232292576788530).members
                        if member not in guild.get_role(1526911990495445143).members
                            and member.id != guild_data["mentor_id"]
                    ]
                    new_mentor = random.choice(members)
                    new_mentor_id = new_mentor.id
                    guild_data["mentor_id"] = new_mentor_id
                    new_next_roll = datetime.datetime(year=today.year,month=today.month,day=today.day) + datetime.timedelta(days=7)
                    guild_data["next_roll"] = new_next_roll.isoformat()
                    LOGGER.debug(f"New mentor: {new_mentor.name} for guild {guild.name}")
        self.save_to_file()

    @commands.Cog.listener()
    async def on_ready(self):
        self.load_from_file()
        self.update.start()

    @commands.Cog.listener()
    async def on_voice_state_update(self, member:discord.Member, before:discord.VoiceState, after:discord.VoiceState):
        if not self.data:
            return

        guild_id = None
        if before and before.channel and before.channel.guild:
            guild_id = before.channel.guild.id
        elif after and after.channel and after.channel.guild:
            guild_id = after.channel.guild.id

        if guild_id is None:
            return

        guild_data = self._get_guild_data(guild_id)
        if not guild_data or member.id != guild_data["mentor_id"]:
            return

        if before and before.channel:
            channel_name = getattr(before.channel, "name", "unknown")
            LOGGER.debug(f"Mentor left channel {channel_name}")

        if after and after.channel:
            channel_name = getattr(after.channel, "name", "unknown")
            LOGGER.debug(f"Mentor joined channel {channel_name}")

        guild_voice_client = self._get_guild_voice_client(guild_id)

        action = self._resolve_voice_action(before.channel if before else None, after.channel if after else None, guild_voice_client)

        if action == "connect":
            try:
                await after.channel.connect()
                self.connected_guild_ids.add(guild_id)
                self.in_voice_channel = True
            except discord.ClientException as exc:
                LOGGER.warning(f"Could not connect to voice channel: {exc}")
        elif action == "move":
            if guild_voice_client:
                await guild_voice_client.move_to(after.channel)
            else:
                await after.channel.connect()
            self.connected_guild_ids.add(guild_id)
            self.in_voice_channel = True
        elif action == "disconnect":
            if guild_voice_client:
                await guild_voice_client.disconnect(force=True)
            self.connected_guild_ids.discard(guild_id)
            self.in_voice_channel = False