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

import re
from typing import Optional, TypedDict

from discord import Guild, Interaction, Message, SelectOption, TextChannel, TextStyle
from discord.app_commands import Choice, Group, autocomplete
from discord.ext.commands import Cog
from discord.ui import Label, Modal, Select, TextInput

from config import CONFIG
from rocketbot.bot import Rocketbot
from rocketbot.cogs.basecog import BaseCog
from rocketbot.cogsetting import CogSetting
from rocketbot.ui.pagedcontent import PAGE_BREAK, paginate, update_paged_content
from rocketbot.utils import (
	MOD_PERMISSIONS,
	blockquote_markdown,
	dump_stacktrace,
	indent_markdown,
)

_CURRENT_DATA_VERSION = 1
_MAX_CONTENT_LENGTH = 2000

class BangCommand(TypedDict):
	content: str
	mod_only: bool
	version: int

async def command_autocomplete(interaction: Interaction, text: str) -> list[Choice[str]]:
	cmds = BangCommandCog.shared.get_saved_commands(interaction.guild)
	return [
		Choice(name=f'!{name}', value=name)
		for name, cmd in sorted(cmds.items())
		if len(text) == 0 or text.lower() in name
	]

class BangCommandCog(BaseCog, name='Bang Commands'):
	SETTING_COMMANDS = CogSetting(
		name='commands',
		datatype=dict[str, BangCommand],
		default_value={},
	)

	shared: Optional['BangCommandCog'] = None

	def __init__(self, bot: Rocketbot):
		super().__init__(
			bot,
			config_prefix='bangcommand',
			short_description='Provides custom informational chat !commands.',
			long_description='Bang commands are simple one-word chat messages starting with an exclamation '
							 '(a "bang") that will make the bot respond with simple informational replies. '
							 'This functionality is similar to Twitch bots. Useful for posting answers to '
							 'frequently asked questions, reminding users of rules, and similar canned responses. '
							 'Commands can be individually made mod-only or usable by anyone.'
		)
		BangCommandCog.shared = self

	def get_saved_commands(self, guild: Guild) -> dict[str, BangCommand]:
		return self.get_guild_setting(guild, BangCommandCog.SETTING_COMMANDS)

	def get_saved_command(self, guild: Guild, name: str) -> Optional[BangCommand]:
		cmds = self.get_saved_commands(guild)
		name = BangCommandCog._normalize_name(name)
		return cmds.get(name, None)

	def set_saved_commands(self, guild: Guild, commands: dict[str, BangCommand]) -> None:
		self.set_guild_setting(guild, BangCommandCog.SETTING_COMMANDS, commands)

	bang = Group(
		name='command',
		description='Provides custom informational chat !commands.',
		guild_only=True,
		default_permissions=MOD_PERMISSIONS,
	)

	@bang.command(
		name='define',
		extras={
			'long_description': 'Simple one-line content can be specified in the command. '
								'For multi-line content, run the command without content '
								'specified to use the editor popup.'
		}
	)
	@autocomplete(name=command_autocomplete)
	async def define_command(self, interaction: Interaction, name: str, definition: Optional[str] = None, mod_only: bool = False) -> None:
		"""
		Defines or redefines a bang command.

		Parameters
		----------
		interaction: Interaction
		name: string
			name of the command (lowercase a-z, underscores, and hyphens)
		definition: string
			content of the command
		mod_only: bool
			whether the command will only be recognized when a mod uses it
		"""
		self.log(interaction.guild, f'{interaction.user.name} used command /bangcommand define {name} {definition} {mod_only}')
		name = BangCommandCog._normalize_name(name)
		if definition is None:
			cmd = self.get_saved_command(interaction.guild, name)
			await interaction.response.send_modal(
				_EditModal(
					name,
					content=cmd['content'] if cmd else None,
					mod_only=cmd['mod_only'] if cmd else None,
					exists=cmd is not None,
				)
			)
			return
		try:
			self.define(interaction.guild, name, definition, mod_only)
			await interaction.response.send_message(
				f'{CONFIG["success_emoji"]} Command `!{name}` has been defined.\n\n{blockquote_markdown(definition)}',
				ephemeral=True,
			)
		except ValueError as e:
			await interaction.response.send_message(
				f'{CONFIG["failure_emoji"]} {e}',
				ephemeral=True,
			)
			return

	@bang.command(
		name='undefine'
	)
	@autocomplete(name=command_autocomplete)
	async def undefine_command(self, interaction: Interaction, name: str) -> None:
		"""
		Removes a bang command.

		Parameters
		----------
		interaction: Interaction
		name: string
			name of the previously defined command
		"""
		try:
			self.undefine(interaction.guild, name)
			await interaction.response.send_message(
				f'{CONFIG["success_emoji"]} Command `!{name}` removed.',
				ephemeral=True,
			)
		except ValueError as e:
			await interaction.response.send_message(
				f'{CONFIG["failure_emoji"]} {e}',
				ephemeral=True,
			)

	@bang.command(
		name='list'
	)
	async def list_command(self, interaction: Interaction) -> None:
		"""
		Lists all defined bang commands.

		Parameters
		----------
		interaction: Interaction
		"""
		cmds = self.get_saved_commands(interaction.guild)
		if cmds is None or len(cmds) == 0:
			await interaction.response.send_message(
				f'{CONFIG["info_emoji"]} No commands defined.',
				ephemeral=True,
				delete_after=15,
			)
			return
		text = '## Commands'
		for name, cmd in sorted(cmds.items()):
			text += PAGE_BREAK + f'\n- `!{name}`'
			if cmd['mod_only']:
				text += ' - **mod only**'
			text += f'\n{indent_markdown(cmd["content"])}'
		pages = paginate(text)
		await update_paged_content(interaction, None, 0, pages)

	@bang.command(
		name='invoke',
		extras={
			'long_description': 'Useful when you do not want the command to show up in chat.',
		}
	)
	@autocomplete(name=command_autocomplete)
	async def invoke_command(self, interaction: Interaction, name: str) -> None:
		"""
		Invokes a bang command without typing it in chat.

		Parameters
		----------
		interaction: Interaction
		name: string
			the bang command name
		"""
		cmd = self.get_saved_command(interaction.guild, name)
		if cmd is None:
			await interaction.response.send_message(
				f'{CONFIG["failure_emoji"]} Command `!{name}` does not exist.',
				ephemeral=True,
			)
			return
		resp = await interaction.response.defer(ephemeral=True, thinking=False)
		await interaction.channel.send(
			cmd['content']
		)
		if resp.resource:
			await resp.resource.delete()

	def define(self, guild: Guild, name: str, content: str, mod_only: bool, check_exists: bool = False) -> None:
		if not BangCommandCog._is_valid_name(name):
			raise ValueError('Invalid command name. Must consist of lowercase letters, underscores, and hyphens (no spaces).')
		name = BangCommandCog._normalize_name(name)
		if len(content) < 1 or len(content) > 2000:
			raise ValueError(f'Content must be between 1 and {_MAX_CONTENT_LENGTH} characters.')
		cmds = self.get_saved_commands(guild)
		if check_exists:
			if cmds.get(name, None) is not None:
				raise ValueError(f'Command with name "{name}" already exists.')
		cmds[name] = {
			'content': content,
			'mod_only': mod_only,
			'version': _CURRENT_DATA_VERSION,
		}
		self.set_saved_commands(guild, cmds)

	def undefine(self, guild: Guild, name: str) -> None:
		name = BangCommandCog._normalize_name(name)
		cmds = self.get_saved_commands(guild)
		if cmds.get(name, None) is None:
			raise ValueError(f'Command with name "{name}" does not exist.')
		del cmds[name]
		self.set_saved_commands(guild, cmds)

	@Cog.listener()
	async def on_message(self, message: Message) -> None:
		if message.guild is None or message.channel is None or not isinstance(message.channel, TextChannel):
			return
		content = message.content
		if content is None or not content.startswith('!') or not BangCommandCog._is_valid_name(content):
			return
		name = BangCommandCog._normalize_name(content)
		cmd = self.get_saved_command(message.guild, name)
		if cmd is None:
			return
		if cmd['mod_only'] and not message.author.guild_permissions.ban_members:
			return
		text = cmd["content"]
		# text = f'{text}\n\n-# {message.author.name} used `!{name}`'
		await message.channel.send(
			text,
		)

	@staticmethod
	def _normalize_name(name: str) -> str:
		name = name.lower().strip()
		if name.startswith('!'):
			name = name[1:]
		return name

	@staticmethod
	def _is_valid_name(name: Optional[str]) -> bool:
		if name is None:
			return False
		return re.match(r'^!?([a-z]+)([_-][a-z]+)*$', name) is not None

class _EditModal(Modal, title='Edit Command'):
	name_label = Label(
		text='Command name',
		description='What gets typed in chat to trigger the command. Must be a-z, underscores, and hyphens (no spaces).',
		component=TextInput(
			style=TextStyle.short, # one line
			placeholder='!command_name',
			min_length=1,
			max_length=100,
		)
	)
	content_label = Label(
		text='Content',
		description='The text the bot will respond with when someone uses the command. Can contain markdown.',
		component=TextInput(
			style=TextStyle.paragraph,
			placeholder='Lorem ipsum dolor...',
			min_length=1,
			max_length=2000,
		)
	)
	mod_only_label = Label(
		text='Mod only?',
		description='Whether mods are the only users who can invoke this command.',
		component=Select(
			options=[
				SelectOption(label='No', value='False',
							 description='Anyone can invoke this command.'),
				SelectOption(label='Yes', value='True',
							 description='Only mods can invoke this command.'),
			],
		)
	)

	def __init__(self, name: Optional[str] = None, content: Optional[str] = None, mod_only: Optional[bool] = None, exists: bool = False):
		super().__init__()
		self.exists = exists
		# noinspection PyTypeChecker
		name_input: TextInput = self.name_label.component
		# noinspection PyTypeChecker
		content_input: TextInput = self.content_label.component
		# noinspection PyTypeChecker
		mod_only_input: Select = self.mod_only_label.component
		name_input.default = name
		content_input.default = content
		resolved_mod_only = mod_only if mod_only is not None else False
		mod_only_input.options[0].default = not resolved_mod_only
		mod_only_input.options[1].default = resolved_mod_only

	async def on_submit(self, interaction: Interaction) -> None:
		# noinspection PyTypeChecker
		name_input: TextInput = self.name_label.component
		# noinspection PyTypeChecker
		content_input: TextInput = self.content_label.component
		# noinspection PyTypeChecker
		mod_only_input: Select = self.mod_only_label.component
		name = name_input.value
		content = content_input.value
		mod_only = mod_only_input.values[0] == 'True'
		try:
			BangCommandCog.shared.define(interaction.guild, name, content, mod_only, not self.exists)
			await interaction.response.send_message(
				f'{CONFIG["success_emoji"]} Command `!{name}` has been defined.\n\n{blockquote_markdown(content)}',
				ephemeral=True,
			)
		except ValueError as e:
			await interaction.response.send_message(
				f'{CONFIG["failure_emoji"]} {e}',
				ephemeral=True,
			)

	async def on_error(self, interaction: Interaction, error: Exception) -> None:
		dump_stacktrace(error)
		try:
			await interaction.response.send_message(
				f'{CONFIG["failure_emoji"]} Save failed',
				ephemeral=True,
			)
		except BaseException:
			pass