From 7cbe1fe198566e2b4ff1d18742da79846a2d255a Mon Sep 17 00:00:00 2001
From: Abhijeet Krishnan <abhijeet.krishnan@gmail.com>
Date: Sat, 17 Feb 2024 16:56:14 -0500
Subject: [PATCH] Add autocomplete functionality

Uses fuzzy string matching from `rapidfuzz`. Still needs work to refactor
the ad-hoc similarity being computed in the FrameDb class.
---
 pyproject.toml         |  3 ++-
 src/framedb/framedb.py |  2 ++
 src/heihachi/bot.py    | 21 +++++++++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7960056..e4e99f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,7 +23,8 @@ dependencies = [
     "discord.py",
     "beautifulsoup4",
     "Requests",
-    "lxml"
+    "lxml",
+    "rapidfuzz"
 ]
 
 [project.urls]
diff --git a/src/framedb/framedb.py b/src/framedb/framedb.py
index ca292ad..59c4513 100644
--- a/src/framedb/framedb.py
+++ b/src/framedb/framedb.py
@@ -83,6 +83,8 @@ def _is_command_in_alias(command: str, move: Move) -> bool:
     def _correct_character_name(char_name_query: str) -> str | None:
         "Check if input in dictionary or in dictionary values"
 
+        char_name_query = char_name_query.lower().strip()
+        char_name_query = char_name_query.replace(" ", "_")
         try:
             return CharacterName(char_name_query).value
         except ValueError:
diff --git a/src/heihachi/bot.py b/src/heihachi/bot.py
index 7cd0d70..febdbe8 100644
--- a/src/heihachi/bot.py
+++ b/src/heihachi/bot.py
@@ -4,6 +4,8 @@
 
 import discord
 import discord.ext.commands
+from rapidfuzz import process
+from rapidfuzz.distance import JaroWinkler
 
 from framedb import CharacterName, FrameDb, FrameService
 from heihachi import button, embed
@@ -75,10 +77,29 @@ async def _character_command(interaction: discord.Interaction["FrameDataBot"], m
 
         return _character_command
 
+    async def _character_name_autocomplete(
+        self, interaction: discord.Interaction["FrameDataBot"], current: str
+    ) -> List[discord.app_commands.Choice[str]]:
+        """
+        Autocomplete function for character names
+
+        Ref.: https://stackoverflow.com/a/75912806/6753162
+        """
+        # TODO: honestly, just build a trie of all char names and aliases and use it here
+
+        char_choices = [char.pretty() for char in CharacterName]
+        trimmed_choices = process.extract(
+            current.lower(), char_choices, scorer=JaroWinkler.distance, limit=3, score_cutoff=0.9
+        )
+        return [discord.app_commands.Choice(name=char, value=char) for char, score, idx in trimmed_choices][
+            :25
+        ]  # Discord has a max choice number of 25 (https://github.com/Rapptz/discord.py/discussions/9241)
+
     def _add_bot_commands(self) -> None:
         "Add all frame commands to the bot"
 
         @self.tree.command(name="fd", description="Frame data from a character move")
+        @discord.app_commands.autocomplete(character=self._character_name_autocomplete)
         async def _frame_data_cmd(
             interaction: discord.Interaction["FrameDataBot"], character: str, move: str
         ) -> None:  # TODO: use command argument completion for char names