python-app-rocketbot/rocketbot/cogs/helpcog.py

import re
from typing import Union, Optional

from discord import Interaction, Permissions, AppCommandType
from discord.app_commands import Group, Command, autocomplete, guild_only, command, Choice

from config import CONFIG
from rocketbot.bot import Rocketbot
from rocketbot.cogs.basecog import BaseCog
from rocketbot.utils import MOD_PERMISSIONS, dump_stacktrace

async def command_autocomplete(interaction: Interaction, current: str) -> list[Choice[str]]:
	"""Autocomplete handler for top-level command names."""
	choices: list[Choice] = []
	try:
		if current.startswith('/'):
			current = current[1:]
		current = current.lower().strip()
		user_permissions = interaction.permissions
		cmds = HelpCog.shared.get_command_list(user_permissions)
		return [
			Choice(name=f'/{cmdname} command', value=f'cmd:{cmdname}')
			for cmdname in sorted(cmds.keys())
			if len(current) == 0 or current in cmdname
		][:25]
	except BaseException as e:
		dump_stacktrace(e)
	return choices

async def subcommand_autocomplete(interaction: Interaction, current: str) -> list[Choice[str]]:
	"""Autocomplete handler for subcommand names. Command taken from previous command token."""
	try:
		current = current.lower().strip()
		cmd_name = interaction.namespace['topic']
		cmd = HelpCog.shared.object_for_help_symbol(cmd_name)
		if not isinstance(cmd, Group):
			return []
		user_permissions = interaction.permissions
		if cmd is None or not isinstance(cmd, Group):
			print(f'No command found named {cmd_name}')
			return []
		grp = cmd
		subcmds = HelpCog.shared.get_subcommand_list(grp, user_permissions)
		if subcmds is None:
			print(f'Subcommands for {cmd_name} was None')
			return []
		return [
			Choice(name=f'{subcmd_name} subcommand', value=f'subcmd:{cmd.name}.{subcmd_name}')
			for subcmd_name in sorted(subcmds.keys())
			if len(current) == 0 or current in subcmd_name
		][:25]
	except BaseException as e:
		dump_stacktrace(e)
	return []

async def cog_autocomplete(interaction: Interaction, current: str) -> list[Choice[str]]:
	"""Autocomplete handler for cog names."""
	try:
		current = current.lower().strip()
		return [
			Choice(name=f'⚙ {cog.qualified_name} module', value=f'cog:{cog.qualified_name}')
			for cog in sorted(HelpCog.shared.bot.cogs.values(), key=lambda c: c.qualified_name)
			if isinstance(cog, BaseCog) and
			   can_use_cog(cog, interaction.permissions) and
			   (len(cog.get_commands()) > 0 or len(cog.settings) > 0) and \
			   (len(current) == 0 or current in cog.qualified_name.lower())
		]
	except BaseException as e:
		dump_stacktrace(e)
	return []

async def topic_autocomplete(interaction: Interaction, current: str) -> list[Choice[str]]:
	"""Autocomplete handler that combines slash commands and cog names."""
	command_choices = await command_autocomplete(interaction, current)
	cog_choices = await cog_autocomplete(interaction, current)
	return (command_choices + cog_choices)[:25]

async def subtopic_autocomplete(interaction: Interaction, current: str) -> list[Choice[str]]:
	"""Autocomplete handler for subtopic names. Currently just handles subcommands."""
	subcommand_choices = await subcommand_autocomplete(interaction, current)
	return subcommand_choices[:25]

class HelpCog(BaseCog, name='Help'):
	shared: Optional['HelpCog'] = None

	def __init__(self, bot: Rocketbot):
		super().__init__(
			bot,
			config_prefix='help',
			short_description='Provides help on using commands and modules.'
		)
		HelpCog.shared = self

	def __create_help_index(self) -> None:
		"""
		Populates self.obj_index and self.keyword_index. Bails if already
		populated. Intended to be run on demand so all cogs and commands have
		had time to get set up and synced.
		"""
		if getattr(self, 'obj_index', None) is not None:
			return
		self.obj_index: dict[str, Union[Command, Group, BaseCog]] = {}
		self.keyword_index: dict[str, set[Union[Command, Group, BaseCog]]] = {}

		def add_text_to_index(obj, text: str):
			words = [
				word
				for word in re.split(r"[^a-zA-Z']+", text.lower())
				if len(word) > 1 and word not in trivial_words
			]
			for word in words:
				matches = self.keyword_index.get(word, set())
				matches.add(obj)
				self.keyword_index[word] = matches

		# PyCharm not interpreting conditional return type correctly.
		# noinspection PyTypeChecker
		cmds: list[Union[Command, Group]] = self.bot.tree.get_commands(type=AppCommandType.chat_input)
		for cmd in cmds:
			key = f'cmd:{cmd.name}'
			self.obj_index[key] = cmd
			add_text_to_index(cmd, cmd.name)
			if cmd.description:
				add_text_to_index(cmd, cmd.description)
			if isinstance(cmd, Group):
				for subcmd in cmd.commands:
					key = f'subcmd:{cmd.name}.{subcmd.name}'
					self.obj_index[key] = subcmd
					add_text_to_index(subcmd, subcmd.name)
					if subcmd.description:
						add_text_to_index(subcmd, subcmd.description)
		for cog_qname, cog in self.bot.cogs.items():
			if not isinstance(cog, BaseCog):
				continue
			key = f'cog:{cog_qname}'
			self.obj_index[key] = cog
			add_text_to_index(cog, cog.qualified_name)
			if cog.description:
				add_text_to_index(cog, cog.description)

	def object_for_help_symbol(self, symbol: str) -> Optional[Union[Command, Group, BaseCog]]:
		self.__create_help_index()
		return self.obj_index.get(symbol, None)

	@command(name='help')
	@guild_only()
	@autocomplete(topic=topic_autocomplete, subtopic=subtopic_autocomplete)
	async def help_command(self, interaction: Interaction, topic: Optional[str] = None, subtopic: Optional[str] = None) -> None:
		"""
		Shows help for using commands and subcommands and configuring modules.

		`/help` will show a list of top-level topics.

		`/help /<command_name>` will show help about a specific command or
		list a command's subcommands.

		`/help /<command_name> <subcommand_name>` will show help about a
		specific subcommand.

		`/help <module_name>` will show help about configuring a module.

		`/help <keywords>` will do a text search for topics.

		Parameters
		----------
		interaction: Interaction
		topic: Optional[str]
			optional command, module, or keywords to get specific help for
		subtopic: Optional[str]
			optional subcommand to get specific help for
		"""
		print(f'help_command(interaction, {topic}, {subtopic})')

		# General help
		if topic is None:
			await self.__send_general_help(interaction)
			return

		# Specific object reference
		obj = self.object_for_help_symbol(subtopic) if subtopic else self.object_for_help_symbol(topic)
		if obj:
			await self.__send_object_help(interaction, obj)
			return

		# Text search
		keywords = [
			word
			for word in re.split(r"[^a-zA-Z']+", topic.lower())
			if len(word) > 0 and word not in trivial_words
		]
		matching_objects_set = None
		for keyword in keywords:
			objs = self.keyword_index.get(keyword, None)
			if objs is not None:
				if matching_objects_set is None:
					matching_objects_set = objs
				else:
					matching_objects_set = matching_objects_set & objs
		accessible_objects = [
			obj
			for obj in matching_objects_set or {}
			if ((isinstance(obj, Command) or isinstance(obj, Group)) and can_use_command(obj, interaction.permissions)) or \
			   (isinstance(obj, BaseCog) and can_use_cog(obj, interaction.permissions))
		]
		await self.__send_keyword_help(interaction, accessible_objects)

	async def __send_object_help(self, interaction: Interaction, obj: Union[Command, Group, BaseCog]) -> None:
		if isinstance(obj, Command):
			if obj.parent:
				await self.__send_subcommand_help(interaction, obj.parent, obj)
			else:
				await self.__send_command_help(interaction, obj)
			return
		if isinstance(obj, Group):
			await self.__send_command_help(interaction, obj)
			return
		if isinstance(obj, BaseCog):
			await self.__send_cog_help(interaction, obj)
			return
		print(f'No help for object {obj}')
		await interaction.response.send_message(
			f'{CONFIG["failure_emoji"]} Failed to get help info.',
			ephemeral=True,
			delete_after=10,
		)

	def get_command_list(self, permissions: Optional[Permissions] = None) -> dict[str, Union[Command, Group]]:
		return { cmd.name: cmd for cmd in self.bot.tree.get_commands() if can_use_command(cmd, permissions) }

	def get_subcommand_list(self, cmd: Group, permissions: Optional[Permissions] = None) -> dict[str, Command]:
		return {
			subcmd.name: subcmd
			for subcmd in cmd.commands
			if can_use_command(subcmd, permissions)
		} if can_use_command(cmd, permissions) else {}

	async def __send_general_help(self, interaction: Interaction) -> None:
		user_permissions: Permissions = interaction.permissions
		all_commands = sorted(self.get_command_list(user_permissions).items())
		all_cog_tuples: list[tuple[str, BaseCog]] = [
			cog_tuple
			for cog_tuple in sorted(self.basecog_map.items())
			if can_use_cog(cog_tuple[1], user_permissions) and \
			   (len(cog_tuple[1].settings) > 0)
		]

		text = f'## :information_source: Help'
		if len(all_commands) + len(all_cog_tuples) == 0:
			text = 'Nothing available for your permissions!'

		if len(all_commands) > 0:
			text += '\n### Commands'
			text += '\nType `/help /commandname` for more information.'
			for cmd_name, cmd in sorted(self.get_command_list(user_permissions).items()):
				text += f'\n- `/{cmd_name}`: {cmd.description}'
				if isinstance(cmd, Group):
					subcommand_count = len(cmd.commands)
					text += f' ({subcommand_count} subcommands)'

		if len(all_cog_tuples) > 0:
			text += '\n### Module Configuration'
			for cog_name, cog in all_cog_tuples:
				has_enabled = next((s for s in cog.settings if s.name == 'enabled'), None) is not None
				text += f'\n- **{cog_name}**: {cog.short_description}'
				if has_enabled:
					text += f'\n   - `/enable {cog.config_prefix}` or `/disable {cog.config_prefix}`'
				for setting in cog.settings:
					if setting.name == 'enabled':
						continue
					text += f'\n   - `/get` or `/set {cog.config_prefix}_{setting.name}`'
		await interaction.response.send_message(
			text,
			ephemeral=True,
		)

	async def __send_keyword_help(self, interaction: Interaction, matching_objects: Optional[list[Union[Command, Group, BaseCog]]]) -> None:
		matching_commands = [
			cmd
			for cmd in matching_objects or []
			if isinstance(cmd, Command) or isinstance(cmd, Group)
		]
		matching_cogs = [
			cog
			for cog in matching_objects or []
			if isinstance(cog, BaseCog)
		]
		if len(matching_commands) + len(matching_cogs) == 0:
			await interaction.response.send_message(
				f'{CONFIG["failure_emoji"]} No available help topics found.',
				ephemeral=True,
				delete_after=10,
			)
			return
		if len(matching_objects) == 1:
			obj = matching_objects[0]
			await self.__send_object_help(interaction, obj)
			return

		text = '## :information_source: Matching Help Topics'
		if len(matching_commands) > 0:
			text += '\n### Commands'
			for cmd in matching_commands:
				if cmd.parent:
					text += f'\n- `/{cmd.parent.name} {cmd.name}`'
				else:
					text += f'\n- `/{cmd.name}`'
		if len(matching_cogs) > 0:
			text += '\n### Cogs'
			for cog in matching_cogs:
				text += f'\n- {cog.qualified_name}'
		await interaction.response.send_message(
			text,
			ephemeral=True,
		)

	async def __send_command_help(self, interaction: Interaction, command_or_group: Union[Command, Group], addendum: Optional[str] = None) -> None:
		if isinstance(command_or_group, Command):
			print(f"Doc:\n{command_or_group.callback.__doc__}")
		text = ''
		if addendum is not None:
			text += addendum + '\n\n'
		text += f'## :information_source: Command Help\n`/{command_or_group.name}`\n\n{command_or_group.description}'
		if isinstance(command_or_group, Group):
			subcmds: dict[str, Command] = self.get_subcommand_list(command_or_group, permissions=interaction.permissions)
			if len(subcmds) > 0:
				text += '\n### Subcommands:'
				for subcmd_name, subcmd in sorted(subcmds.items()):
					text += f'\n- `{subcmd_name}`: {subcmd.description}'
		else:
			params = command_or_group.parameters
			if len(params) > 0:
				text += '\n### Parameters:'
				for param in params:
					text += f'\n- `{param.name}`: {param.description}'
		await interaction.response.send_message(text, ephemeral=True)

	async def __send_subcommand_help(self, interaction: Interaction, group: Group, subcommand: Command) -> None:
		text = f'## :information_source: Subcommand Help'
		text += f'\n`/{group.name} {subcommand.name}`'
		text += f'\n\n{subcommand.description}'
		params = subcommand.parameters
		if len(params) > 0:
			text += '\n### Parameters:'
			for param in params:
				text += f'\n- `{param.name}`: {param.description}'
		await interaction.response.send_message(text, ephemeral=True)

	async def __send_cog_help(self, interaction: Interaction, cog: BaseCog) -> None:
		text = f'## :information_source: Module Help'
		text += f'\n**{cog.qualified_name}** module'
		if cog.description is not None:
			text += f'\n\n{cog.description}'

		cmds = [
			cmd
			for cmd in sorted(cog.get_app_commands(), key=lambda c: c.name)
			if can_use_command(cmd, interaction.permissions)
		]
		if len(cmds) > 0:
			text += '\n### Commands:'
			for cmd in cmds:
				text += f'\n- `/{cmd.name}` - {cmd.description}'
				if isinstance(cmd, Group):
					subcmds = [ subcmd for subcmd in cmd.commands if can_use_command(subcmd, interaction.permissions) ]
					if len(subcmds) > 0:
						text += f' ({len(subcmds)} subcommands)'

		settings = cog.settings
		if len(settings) > 0:
			text += '\n### Configuration'
			enabled_setting = next((s for s in settings if s.name == 'enabled'), None)
			if enabled_setting is not None:
				text += f'\n- `/enable {cog.config_prefix}` or `/disable {cog.config_prefix}`'
			for setting in sorted(settings, key=lambda s: s.name):
				if setting.name == 'enabled':
					continue
				text += f'\n- `/get` or `/set {cog.config_prefix}_{setting.name}` - {setting.brief}'
		print(text)
		await interaction.response.send_message(
			text,
			ephemeral=True,
		)

# Exclusions from keyword indexing
trivial_words = {
	'a', 'an', 'and', 'are', "aren't", 'as', 'by', 'can', 'for', 'have', 'if', 'in',
	'is', 'it', 'its', "it's", 'not', 'of', 'on', 'or', 'than', 'that', 'the', 'then',
	'there', 'them', 'they', "they're", 'this', 'to', 'when', 'with',
}

def can_use_command(cmd: Union[Group, Command], user_permissions: Optional[Permissions]) -> bool:
	if user_permissions is None:
		return False
	if cmd.parent and not can_use_command(cmd.parent, user_permissions):
		return False
	return cmd.default_permissions is None or cmd.default_permissions.is_subset(user_permissions)

def can_use_cog(cog: BaseCog, user_permissions: Optional[Permissions]) -> bool:
	# "Using" a cog for now means configuring it, and only mods can configure cogs.
	return user_permissions is not None and MOD_PERMISSIONS.is_subset(user_permissions)