(WIP) Refactor heihachi bot code

src/framedb/const.py

@@ -136,23 +136,6 @@ class MoveType(enum.Enum):
     "fc.": "fc",
 }
 
-BLACKLIST = [
-    "ImVeryBad#5231",
-    "Nape Brasslers#1131",
-    "Sleepii#1337",
-    "iaa ibb beb ib#0000",
-    "ГЕНИЙ#5448",
-    "Beeno#6524",
-    "Gigass-7#6960",
-    "nickname#0000",
-    "Sleepii#6666",
-    "scrotum buster#6919",
-    "Woozle#6308",
-    "Iam#1001",
-    "ImVeryBad#5231",
-]
-ID_BLACKLIST = [1006234003893915679, 240521686531702784]
-
 EMOJI_LIST = [
     "1\ufe0f\u20e3",
     "2\ufe0f\u20e3",

src/framedb/framedb.py

@@ -5,7 +5,7 @@
 from typing import Dict, List
 
 from .character import Character, Move
-from .const import MOVE_TYPE_ALIAS, REPLACE, CharacterName, MoveType
+from .const import CHARACTER_ALIAS, MOVE_TYPE_ALIAS, REPLACE, CharacterName, MoveType
 from .frame_service import FrameService
 
 logger = logging.getLogger(__name__)
@@ -25,13 +25,19 @@ def export(self, export_dir_path: str, format: str = "json") -> None:
 
         if not os.path.exists(export_dir_path):
             os.makedirs(export_dir_path)
-        pass  # TODO: implement this
+
+        match format:
+            case "json":
+                for character in self.frames.values():
+                    character.export_movelist(os.path.join(export_dir_path, f"{character.name}.{format}"), format=format)
+                    logger.info(f"Exported frame data for {character.name} to {export_dir_path}/{character.name}.{format}")
 
     def load(self, frame_service: FrameService) -> None:
         "Load the frame database using a frame service."
 
         for character in CharacterName:
             frames = frame_service.get_frame_data(character)
+            logger.info(f"Retrieved frame data for {character.value} from {frame_service.name}")
             if frames:
                 self.frames[character] = frames
             else:
@@ -62,6 +68,19 @@ def _is_command_in_alias(command: str, move: Move) -> bool:
                 return True
         return False
 
+    @staticmethod
+    def _correct_character_name(char_name_query: str) -> str | None:
+        "Check if input in dictionary or in dictionary values"
+
+        if char_name_query in CHARACTER_ALIAS:
+            return char_name_query
+
+        for key, value in CHARACTER_ALIAS.items():
+            if char_name_query in value:
+                return key.value
+
+        return None
+
     @staticmethod
     def _correct_move_type(move_type_query: str) -> MoveType | None:
         """Given a move type query, correct it to a known move type."""
@@ -132,6 +151,21 @@ def get_similar_moves(self, character: CharacterName, input_query: str) -> List[
 
         return result
 
+    def get_character_by_name(self, name_query: str) -> Character | None:
+        """Given a character name query, return the corresponding character"""
+
+        for character_name, character in self.frames.items():
+            if character_name.value == name_query:
+                return character
+        return None
+
+    def get_move_type(self, move_type_query: str) -> MoveType | None:
+        """Given a move type query, return the corresponding move type"""
+
+        for move_type, aliases in MOVE_TYPE_ALIAS.items():
+            if move_type_query.lower() in aliases:
+                return move_type
+
 
 def _get_close_matches_indices(word: str, possibilities: List[str], n: int = 5, cutoff: float = 0.7) -> List[int]:
     """

src/heihachi/__init__.py

@@ -0,0 +1 @@
+from .configurator import Configurator

src/heihachi/bot.py

@@ -0,0 +1,111 @@
+import datetime
+import logging
+import sched
+
+import discord
+
+from framedb.const import CharacterName
+from heihachi import Configurator, button, embed
+from heihachi.embed import create_frame_data_embed
+
+logger = logging.getLogger(__name__)
+
+
+class FrameDataBot(discord.Client):
+    def __init__(self, config: Configurator, intents: discord.Intents):
+        super().__init__(intents=intents)
+        self.config = config
+        self.synced = False
+
+    async def on_ready(self, tree: discord.app_commands.CommandTree) -> None:
+        await self.wait_until_ready()
+        if not self.synced:
+            await tree.sync()
+            self.synced = True
+        logger.info(f"Logged on as {self.user}")
+
+    def is_user_blacklisted(self, user_id: str | int) -> bool:
+        "Check if a user is blacklisted"
+
+        if isinstance(user_id, str):
+            blacklist = self.config.blacklist
+        else:
+            blacklist = self.config.id_blacklist
+
+        if blacklist:
+            return user_id in blacklist
+        else:
+            return False
+
+    def is_author_newly_created(self, interaction: discord.Interaction) -> bool:
+        "Check if author of an interaction is newly created"
+
+        today = datetime.datetime.strptime(datetime.datetime.now().isoformat(), "%Y-%m-%dT%H:%M:%S.%f")
+        age = today - interaction.user.created_at.replace(tzinfo=None)
+        return age.days < self.config.new_author_age_limit
+
+
+# TODO: fix all this
+@hei.event
+async def on_message(message) -> None:
+    if not self.is_user_blacklisted(message.author.id) and message.content and message.author.id != hei.user.id:
+        user_command = message.content.split(" ", 1)[1]
+        parameters = user_command.strip().split(" ", 1)
+        character_name = parameters[0].lower()
+        character_move = parameters[1]
+
+        embed = create_frame_data_embed(character_name, character_move)
+        await message.channel.send(embed=embed)
+
+
+@tree.command(name="fd", description="Frame data from a character move")
+async def self(interaction: discord.Interaction, character_name: str, move: str) -> None:
+    if not (self.is_user_blacklisted(str(interaction.user.id)) or self.is_author_newly_created(interaction)):
+        embed = create_frame_data_embed(character_name, move)
+        await interaction.response.send_message(embed=embed, ephemeral=False)
+
+
+def character_command_factory(name: str):
+    async def command(interaction: discord.Interaction, move: str) -> None:
+        if not (self.is_user_blacklisted(str(interaction.user.id)) or self.is_author_newly_created(interaction)):
+            embed = create_frame_data_embed(name, move)
+            await interaction.response.send_message(embed=embed, ephemeral=False)
+
+    return command
+
+
+for character in CharacterName:
+    name = character.value
+    tree.command(name=name, description=f"Frame data from {name}")(character_command_factory(name))
+
+if self.config.feedback_channel_id:
+
+    @tree.command(name="feedback", description="Send feedback incase of wrong data")
+    async def self(interaction: discord.Interaction, message: str) -> None:
+        if not (self.is_user_blacklisted(str(interaction.user.id)) or util.is_author_newly_created(interaction)):
+            try:
+                feedback_message = "Feedback from **{}** with ID **{}** in **{}** \n- {}\n".format(
+                    str(interaction.user.name),
+                    interaction.user.id,
+                    interaction.guild,
+                    message,
+                )
+                try:
+                    channel = hei.get_channel(config.feedback_channel_id)
+                    actioned_channel = hei.get_channel(config.action_channel_id)
+                except Exception as e:
+                    logger.error(f"Error getting channel: {e}")
+                await channel.send(content=feedback_message, view=button.DoneButton(actioned_channel))
+                result = embed.success_embed("Feedback sent")
+            except Exception as e:
+                result = embed.error_embed(f"Feedback couldn't be sent, caused by: {str(e)}")
+
+            await interaction.response.send_message(embed=result, ephemeral=False)
+else:
+    logger.warning("Feedback channel ID is not set. Disabling feedback command.")
+
+
+def periodic_function(scheduler: sched.scheduler, interval: float, function: sched._ActionCallback, character_list_path: str):
+    while True:
+        scheduler.enter(interval, 1, function, (character_list_path,))
+        scheduler.run()

src/heihachi/configurator.py

@@ -1,7 +1,7 @@
 import json
 import logging
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any, Dict, List, Optional
 
 logger = logging.getLogger(__name__)
 
@@ -15,12 +15,18 @@ class Configurator:
     discord_token: str
     feedback_channel_id: int | None
     action_channel_id: int | None
+    blacklist: List[str] | None
+    id_blacklist: List[int] | None
+    new_author_age_limit: int = 120
 
-    def to_dict(self) -> dict[str, Any]:
+    def to_dict(self) -> Dict[str, Any]:
         return {
             "DISCORD_TOKEN": self.discord_token,
             "FEEDBACK_CHANNEL_ID": self.feedback_channel_id,
             "ACTION_CHANNEL_ID": self.action_channel_id,
+            "BLACKLIST": self.blacklist,
+            "ID_BLACKLIST": self.id_blacklist,
+            "NEW_AUTHOR_AGE_LIMIT": self.new_author_age_limit,
         }
 
     @staticmethod
@@ -38,6 +44,9 @@ def from_file(config_path: str) -> Optional["Configurator"]:
                 discord_token=config_data["DISCORD_TOKEN"],
                 feedback_channel_id=config_data.get("FEEDBACK_CHANNEL_ID", None),
                 action_channel_id=config_data.get("ACTION_CHANNEL_ID", None),
+                blacklist=config_data.get("BLACKLIST", None),
+                id_blacklist=config_data.get("ID_BLACKLIST", None),
+                new_author_age_limit=config_data.get("NEW_AUTHOR_AGE_LIMIT", 120),
             )
         except FileNotFoundError:
             logger.error(f"Config file not found at {config_path}")

src/heihachi/embed.py

@@ -2,8 +2,7 @@
 
 import discord
 
-from framedb.character import Character, Move
-from framedb.const import CharacterName, MoveType
+from framedb import Character, CharacterName, FrameService, Move, MoveType
 
 MOVE_NOT_FOUND_TITLE = "Move not found"
 
@@ -12,7 +11,9 @@
 ERROR_COLOR = discord.Colour.from_rgb(220, 20, 60)
 
 
-def similar_moves_embed(similar_moves: List[Move], character_name: CharacterName) -> discord.Embed:
+def similar_moves_embed(
+    frame_service: FrameService, similar_moves: List[Move], character_name: CharacterName
+) -> discord.Embed:
     """Returns the embed message for similar moves."""
 
     command_list = [f"**{idx + 1}**. {move.input}" for idx, move in enumerate(similar_moves)]
@@ -25,7 +26,9 @@ def similar_moves_embed(similar_moves: List[Move], character_name: CharacterName
     return embed
 
 
-def move_list_embed(character: Character, moves: List[Move], move_type: MoveType) -> discord.Embed:
+def move_list_embed(
+    frame_service: FrameService, character: Character, moves: List[Move], move_type: MoveType
+) -> discord.Embed:
     """Returns the embed message for a list of moves matching to a special move type."""
 
     desc_string = "\n".join(sorted([move.input for move in moves]))
@@ -48,18 +51,18 @@ def success_embed(message) -> discord.Embed:
     return embed
 
 
-def move_embed(character: Character, move: Move) -> discord.Embed:
+def move_embed(frame_service: FrameService, character: Character, move: Move) -> discord.Embed:
     """Returns the embed message for character and move."""
 
     embed = discord.Embed(
         title=f"**{move.input}**",
         colour=SUCCESS_COLOR,
         description=move.name,
-        url=f"{character.page}_movelist#{move.id}",
+        url=f"{character.page}_movelist#{move.id}",  # TODO: this is specific to Wavu, change it to be more generic
     )
 
     embed.set_thumbnail(url=character.portrait)
-    embed.set_footer(text="Wavu Wiki", icon_url=WAVU_LOGO)
+    embed.set_footer(text=frame_service.name, icon_url=frame_service.icon)
     embed.set_author(name=character.name.value.title(), url=character.page)
 
     embed.add_field(name="Target", value=move.target)
@@ -73,3 +76,30 @@ def move_embed(character: Character, move: Move) -> discord.Embed:
         embed.add_field(name="Notes", value=move.notes)
 
     return embed
+
+
+def create_frame_data_embed(name: str, move: str) -> discord.Embed:
+    character_name = util.correct_character_name(name.lower())  # TODO: fix all this
+    if character_name:
+        character = util.get_character_by_name(character_name, character_list)
+        assert character is not None
+        move_list = json_directory.get_movelist(character_name, JSON_PATH)
+        move_type = util.get_move_type(move)
+
+        if move_type:
+            moves = json_directory.get_by_move_type(move_type, move_list)
+            moves_embed = embed.move_list_embed(character, moves, move_type)
+            return moves_embed
+        else:
+            character_move = json_directory.get_move(move, move_list)
+
+            if character_move:
+                move_embed = embed.move_embed(character, character_move)
+                return move_embed
+            else:
+                similar_moves = json_directory.get_similar_moves(move, move_list)
+                similar_moves_embed = embed.similar_moves_embed(similar_moves, character_name)
+                return similar_moves_embed
+    else:
+        error_embed = embed.error_embed(f"Could not locate character {name}.")
+        return error_embed

src/heihachi/tests/test_bot.py

@@ -0,0 +1,3 @@
+import pytest
+
+# TODO: Add tests for the bot

src/heihachi/tests/test_embed.py

@@ -0,0 +1,26 @@
+import pytest
+
+
+@pytest.mark.skip(reason="Not implemented yet.")
+def test_similar_moves_embed():
+    pass
+
+
+@pytest.mark.skip(reason="Not implemented yet.")
+def test_move_list_embed():
+    pass
+
+
+@pytest.mark.skip(reason="Not implemented yet.")
+def test_error_embed():
+    pass
+
+
+@pytest.mark.skip(reason="Not implemented yet.")
+def test_success_embed():
+    pass
+
+
+@pytest.mark.skip(reason="Not implemented yet.")
+def test_move_embed():
+    pass