Refactor main function into ChaosBot class for improved structure and initialization of components

botbotbot/__init__.py

@@ -1,5 +1,6 @@
 import logging
 import tomllib
+from typing import Any
 
 import discord
 
@@ -11,66 +12,98 @@ from botbotbot.tts import CambAI
 from botbotbot.voice import VoiceBot
 from botbotbot.wordlist import Wordlist
 
+logger = logging.getLogger(__name__)
+
+
+class ChaosBot:
+    def __init__(self, config: dict[str, Any]) -> None:
+        self.config = config
+
+        intents = discord.Intents.all()
+        intents.presences = False
+
+        self.bot = discord.Bot(description="Discord Chaos Bot", intents=intents)
+
+        self.init_aibot()
+        self.init_cambai()
+        self.init_guild_ids()
+
+        self.shuffler = Shuffler()
+        self.word_list = Wordlist(self.bot, self.guild_ids)
+        self.text_bot = TextBot(
+            self.bot,
+            self.word_list,
+            aibot=self.ai_bot,
+            shuffler=self.shuffler,
+            guild_ids=self.guild_ids,
+        )
+        self.voice_bot = VoiceBot(
+            self.bot,
+            self.cambai,
+            aibot=self.ai_bot,
+            shuffler=self.shuffler,
+            guild_ids=self.guild_ids,
+        )
+        self.nick_bot = NickBot(
+            self.bot, shuffler=self.shuffler, guild_ids=self.guild_ids
+        )
+
+        self.init_events()
+
+    def init_aibot(self) -> None:
+        self.ai_bot: AIBot | None = None
+
+        if isinstance(key := self.config.get("mistral_api_key"), str):
+            system_prompt = """Tu es une intelligence artificelle qui répond en français.
+            Tu dois faire un commentaire pertinent en lien avec ce qui te sera dit.
+            Ta réponse doit être très courte.
+            Ta réponse doit être une seule phrase.
+            TA RÉPONSE DOIT ÊTRE EN FRANÇAIS !!!"""
+
+            self.ai_bot = AIBot(
+                key,
+                model="mistral-large-latest",
+                system_message=system_prompt,
+            )
+        else:
+            logger.warning("No AI Bot.")
+
+    def init_cambai(self) -> None:
+        self.cambai: CambAI | None = None
+        if isinstance(key := self.config.get("cambai_api_key"), str):
+            self.cambai = CambAI(key)
+        else:
+            logger.warning("No CambAI.")
+
+    def init_guild_ids(self) -> None:
+        guild_ids = self.config.get("guild_ids")
+        if not (
+            isinstance(guild_ids, list) and all(isinstance(i, int) for i in guild_ids)
+        ):
+            logger.error("Guild IDs must be a list of integers.")
+            guild_ids = []
+        self.guild_ids = guild_ids
+
+    def init_events(self) -> None:
+        self.bot.add_listener(self.on_ready, "on_ready")
+
+        self.word_list.init_events()
+        self.text_bot.init_events()
+        self.voice_bot.init_events()
+        self.nick_bot.init_events()
+
+    async def on_ready(self) -> None:
+        logger.info(f"We have logged in as {self.bot.user}")
+
+    def run(self) -> None:
+        self.bot.run(self.config.get("token"))
+
 
 def main() -> None:
-    description = """Discord Chaos Bot"""
-
-    logger = logging.getLogger(__name__)
     logging.basicConfig(level=logging.INFO)
 
     with open("config.toml", "rb") as config_file:
         config = tomllib.load(config_file)
 
-    guild_ids = config.get("guild_ids")
+    chaos_bot = ChaosBot(config)
-    if not (isinstance(guild_ids, list) and all(isinstance(i, int) for i in guild_ids)):
+    chaos_bot.run()
-        logger.error("Guild IDs must be a list of integers.")
-        guild_ids = []
-
-    system_prompt = """Tu es une intelligence artificelle qui répond en français.
-    Tu dois faire un commentaire pertinent en lien avec ce qui te sera dit.
-    Ta réponse doit être très courte.
-    Ta réponse doit être une seule phrase.
-    TA RÉPONSE DOIT ÊTRE EN FRANÇAIS !!!"""
-
-    aibot: AIBot | None = None
-
-    if isinstance(key := config.get("mistral_api_key"), str):
-        aibot = AIBot(
-            key,
-            model="mistral-large-latest",
-            system_message=system_prompt,
-        )
-    else:
-        logger.warning("No AI Bot.")
-
-    cambai: CambAI | None = None
-    if isinstance(key := config.get("cambai_api_key"), str):
-        cambai = CambAI(key)
-    else:
-        logger.warning("No CambAI.")
-
-    intents = discord.Intents.default()
-    intents.members = True
-    intents.message_content = True
-    intents.reactions = True
-    intents.voice_states = True
-
-    bot = discord.Bot(description=description, intents=intents)
-    shf = Shuffler()
-    wl = Wordlist(bot, guild_ids)
-    wl.init_events()
-
-    text_bot = TextBot(bot, wl, aibot=aibot, shuffler=shf, guild_ids=guild_ids)
-    text_bot.init_events()
-
-    voice_bot = VoiceBot(bot, cambai, aibot=aibot, shuffler=shf, guild_ids=guild_ids)
-    voice_bot.init_events()
-
-    nick_bot = NickBot(bot, shuffler=shf, guild_ids=guild_ids)
-    nick_bot.init_events()
-
-    @bot.listen("on_ready")
-    async def on_ready() -> None:
-        logger.info(f"We have logged in as {bot.user}")
-
-    bot.run(config.get("token"))