diff --git a/flanabot/bots/flana_disc_bot.py b/flanabot/bots/flana_disc_bot.py index 331761b..7bce741 100644 --- a/flanabot/bots/flana_disc_bot.py +++ b/flanabot/bots/flana_disc_bot.py @@ -8,7 +8,6 @@ import random from collections import defaultdict import discord -import flanautils import pytz from flanautils import Media, NotFoundError, OrderedSet from multibot import BadRoleError, DiscordBot, Platform, Role, User, admin, bot_mentioned, constants as multibot_constants, group @@ -16,6 +15,7 @@ from multibot import BadRoleError, DiscordBot, Platform, Role, User, admin, bot_ from flanabot import constants from flanabot.bots.flana_bot import FlanaBot from flanabot.models import Chat, Message, Punishment +from models.heating_context import ChannelData, HeatingContext # ---------------------------------------------------------------------------------------------------- # @@ -24,9 +24,7 @@ from flanabot.models import Chat, Message, Punishment class FlanaDiscBot(DiscordBot, FlanaBot): def __init__(self): super().__init__(os.environ['DISCORD_BOT_TOKEN']) - self.channels_data: dict[int, dict[str, dict]] = defaultdict(dict) - self.heating = False - self.heat_level = 0.0 + self.heating_contexts: dict[int, HeatingContext] = defaultdict(HeatingContext) # -------------------------------------------------------- # # ------------------- PROTECTED METHODS ------------------ # @@ -47,51 +45,55 @@ class FlanaDiscBot(DiscordBot, FlanaBot): async def _heat_channel(self, channel: discord.VoiceChannel): async def set_fire_to(channel_key: str, depends_on: str, firewall=0): - fire_score = random.randint(0, guild_channels_data[depends_on]['n_fires'] - guild_channels_data[channel_key]['n_fires']) - firewall // 2 + fire_score = random.randint(0, channels_data[depends_on].n_fires - channels_data[channel_key].n_fires) - firewall // 2 if fire_score < 1: - if not guild_channels_data[channel_key]['n_fires']: + if not channels_data[channel_key].n_fires: return - guild_channels_data[channel_key]['n_fires'] -= 1 + channels_data[channel_key].n_fires -= 1 elif fire_score == 1: return else: - guild_channels_data[channel_key]['n_fires'] += 1 + channels_data[channel_key].n_fires += 1 - if guild_channels_data[channel_key]['n_fires']: - new_name_ = '🔥' * guild_channels_data[channel_key]['n_fires'] + if channels_data[channel_key].n_fires: + new_name_ = '🔥' * channels_data[channel_key].n_fires else: - new_name_ = guild_channels_data[channel_key]['original_name'] - await guild_channels_data[channel_key]['object'].edit(name=new_name_) + new_name_ = channels_data[channel_key].original_name + await channels_data[channel_key].channel.edit(name=new_name_) - guild_channels_data = self.channels_data[channel.guild.id] + voice_channels = {} + for voice_channel in channel.guild.voice_channels: + voice_channels[voice_channel.id] = voice_channel + channels_data = {} for letter, channel_id in constants.DISCORD_HOT_CHANNEL_IDS.items(): - channel_ = flanautils.find(channel.guild.voice_channels, condition=lambda c: c.id == channel_id) - guild_channels_data[letter] = { - 'object': channel_, - 'original_name': channel_.name, - 'n_fires': 0 - } + channels_data[letter] = ChannelData( + channel=voice_channels[channel_id], + original_name=voice_channels[channel_id].name + ) + + heating_context = self.heating_contexts[channel.guild.id] + heating_context.channels_data = channels_data while True: - await asyncio.sleep(constants.HEAT_PERIOD_SECONDS) + await asyncio.sleep(4) # constants.HEAT_PERIOD_SECONDS) if channel.members: - self.heat_level += 0.5 + heating_context.heat_level += 0.5 else: - if not self.heat_level: + if not heating_context.heat_level: return - self.heat_level -= 0.5 - if self.heat_level > len(constants.DISCORD_HEAT_NAMES) - 1: - self.heat_level = float(int(self.heat_level)) + heating_context.heat_level -= 0.5 + if heating_context.heat_level > len(constants.DISCORD_HEAT_NAMES) - 1: + heating_context.heat_level = float(int(heating_context.heat_level)) - if not self.heat_level.is_integer(): + if not heating_context.heat_level.is_integer(): continue - i = int(self.heat_level) + i = int(heating_context.heat_level) if not i: n_fires = 0 - new_name = guild_channels_data['C']['original_name'] + new_name = channels_data['C'].original_name elif i < len(constants.DISCORD_HEAT_NAMES): n_fires = 0 new_name = constants.DISCORD_HEAT_NAMES[i] @@ -99,14 +101,14 @@ class FlanaDiscBot(DiscordBot, FlanaBot): n_fires = i - len(constants.DISCORD_HEAT_NAMES) + 1 n_fires = round(math.log(n_fires + 4, 1.2) - 8) new_name = '🔥' * n_fires - guild_channels_data['C']['n_fires'] = n_fires + channels_data['C'].n_fires = n_fires if channel.name != new_name: await channel.edit(name=new_name) - await set_fire_to('B', depends_on='C', firewall=len(guild_channels_data['B']['object'].members)) - await set_fire_to('A', depends_on='B', firewall=len(guild_channels_data['A']['object'].members)) - await set_fire_to('D', depends_on='C', firewall=len(guild_channels_data['C']['object'].members)) - await set_fire_to('E', depends_on='D', firewall=len(guild_channels_data['D']['object'].members)) + await set_fire_to('B', depends_on='C', firewall=len(channels_data['B'].channel.members)) + await set_fire_to('A', depends_on='B', firewall=len(channels_data['A'].channel.members)) + await set_fire_to('D', depends_on='C', firewall=len(channels_data['C'].channel.members)) + await set_fire_to('E', depends_on='D', firewall=len(channels_data['D'].channel.members)) async def _punish(self, user: int | str | User, group_: int | str | Chat | Message, message: Message = None): user_id = self.get_user_id(user) @@ -188,10 +190,11 @@ class FlanaDiscBot(DiscordBot, FlanaBot): else: return - if not self.heating: - self.heating = True + heating_context = self.heating_contexts[channel.guild.id] + if not heating_context.is_active: + heating_context.is_active = True await self._heat_channel(channel) - self.heating = False + heating_context.is_active = False # -------------------------------------------------------- # # -------------------- PUBLIC METHODS -------------------- # @@ -208,11 +211,11 @@ class FlanaDiscBot(DiscordBot, FlanaBot): })) async def restore_channel_names(self, group_id: int): - for channel_data in self.channels_data[group_id].values(): - original_name = channel_data['original_name'] - channel = channel_data['object'] - if channel.name != original_name: - await channel.edit(name=original_name) - channel_data['n_fires'] = 0 + heating_context = self.heating_contexts[group_id] - self.heat_level = 0.0 + for channel_data in heating_context.channels_data.values(): + if channel_data.channel.name != channel_data.original_name: + await channel_data.channel.edit(name=channel_data.original_name) + channel_data.n_fires = 0 + + heating_context.heat_level = 0.0 diff --git a/flanabot/models/heating_context.py b/flanabot/models/heating_context.py new file mode 100644 index 0000000..0e4e48a --- /dev/null +++ b/flanabot/models/heating_context.py @@ -0,0 +1,19 @@ +__all__ = ['ChannelData', 'HeatingContext'] + +from dataclasses import dataclass, field + +import discord + + +@dataclass +class ChannelData: + channel: discord.VoiceChannel + original_name: str + n_fires: int = 0 + + +@dataclass +class HeatingContext: + channels_data: dict[str, ChannelData] = field(default_factory=dict) + is_active: bool = False + heat_level: float = 0.0