umi/cogs/random.py

import discord
from discord.ext.commands import Cog, command, Context
from discord.commands import slash_command, Option
import httpx
from decouple import config
from collections import Counter
import re
from typing import Iterator

GUILD = config("DISCORD_GUILD_ID", cast=int)

def setup(bot: discord.Bot):
	bot.add_cog(Random(bot))

class Random(Cog):
	"""Use random.org to provide truly random results"""
	def __init__(self, bot: discord.Bot):
		self.bot: discord.Bot = bot
		print("Initialized Random cog")

# ==============================================================================

	@staticmethod
	def roll(
		number: int = 1,
		faces: int = 6,
		) -> str:
		"""
		Roll standard dice using random.org integer generation
		
		Params:
			int number    How many dice to roll
			int faces     How many faces on a standard die
		"""
		columns = number if number < 11 else 10
		return httpx.get(f'https://www.random.org/integers/?num={number}&min=1&max={faces}&col={columns}&base=10&format=plain&rnd=new').text

	@staticmethod
	def parse_dice_notation(input: str) -> Iterator[re.Match[str]]:
		"""
		Parse dice notation to get numbers, faces, modifiers, and tokens.
		https://en.wikipedia.org/wiki/Dice_notation

		Dice notation can take the form of AdX+B (chainable):

		A             how many dice to roll (if not provided, roll once)
		dX            die with X faces (1-X)
		+             add the result to the previous result
		-             subtract the result from the previous result
		B             constant modifier

		Extensions implemented:

		d(a,b,...)    custom die with faces defined by a sequence
		d%            special die with 00-99
		dF            fudge dice (two +, two -, two blanks)
		dF.1          fudge dice variant (one +, one -, four blanks)
		dF.2          alias for dF

		Unimplemented:

		-L            subtract the lowest
		-H            subtract the highest
		klL           keep lowest L values
		khH           keep highest H values
		xC            multiply the result by C
		Cx(AdX+B)     multiply the dice group by C (including modifier)
		/C            divide by C?

		Params:
			str input              some dice notation to parse

		Returns:
			tuple(plus_minus_dice, number, faces, plus_minus_mod, modifier)

			str plus_minus_dice    a token for adding or removing from total
			str number             the number of dice to roll
			str faces              the faces on the dice (sequence or number)
			---
			str plus_minus_mod     a token for adding or subtracting a modifier
			str modifier           some number to add/subtract from total
		"""
		pattern = re.compile(
			r'''
			# dice group
			(\A|\+|\-)  # start-of-string or plus or minus dice result (group 1)
			(\d+)?  # number of dice (group 2)
			d  # the literal 'd'
			(  # faces (group 3)
				\d+  # either plain number
				|\(.+?\)  # or sequence of strings
				|\[\d+:\d+(?:\d+)?\]  # or range of numbers as start:end:step
				|%  # or literal '%' (for 00-99)
				|F(?:\.\d+?)?  # or literal 'F' (for fudge dice ++--__)
			)
			# OR
			|
			# number modifier
			(\+|\-)  # plus or minus (group 4)
			(\d+)?  # some number value (group 5)
			''',
			re.IGNORECASE | re.VERBOSE
		)
		# returns (plus_minus_dice, number, faces, plus_minus_mod, modifier)
		return pattern.finditer(input)

	@staticmethod
	def roll_result(roll: str) -> str:
		"""Calculates the result of a roll and output a summary"""
		matches = Random.parse_dice_notation(roll)

		total = 0
		roll_results = []
		for match in matches:
			plus_minus_dice, number, faces, plus_minus_mod, modifier = match.groups()

			if faces and faces[0] == '(':  # non-standard dice sequence (a,b,c,...).
				number = int(number) if number else 1  # number may be omitted for single die.
				outcomes = faces[1:-1].split(',')  # remove parens and get list of faces,
				result = Random.roll(number, len(outcomes))  # then roll for an index in that list
				result = [outcomes[int(r)-1] for r in result.split()]  # that we can then lookup.
				try:  # we don't actually know if the faces are numeric...
					subtotal = sum(map(int, result))
					if plus_minus_dice == '-':
						total -= subtotal
					else:
						total += subtotal
				except (ValueError, TypeError):  # in case faces were not all numeric (for this group or a previous group),
					total = None  # there is no total.
				result = ' '.join(result)  # finally, repack the result
				roll_results.append( (f'{number}d{faces}', result) )  # and append it to output
			elif faces and faces == '%':  # special 00-99
				number = int(number) if number else 1  # number may be omitted for single die.
				result = Random.roll(number, 100)  # we roll a d100,
				result = re.sub('100', '00', result)  # but we sub in 00
				result = re.sub(r'\b(\d){1}\b', r'0\1', result) # and zero pad single digits.
				try:  # this may be part of a larger group of rolls...
					subtotal = sum(map(int, result.split()))
					if plus_minus_dice == '-':
						total -= int(subtotal)  # this may fail because total is None
					else:
						total += int(subtotal)  # this may fail because total is None
				except (ValueError, TypeError):
					total = None
				roll_results.append( (f'{number}d{faces}', result) )  # and append it to output
			elif faces and faces[0] == 'F':  # fudge dice
				number = int(number) if number else 1  # number may be omitted for single die.
				if faces == 'F' or faces == 'F.2':
					outcomes = '(+,+,-,-,␢,␢)'
				elif faces == 'F.1':
					outcomes = '(+,+,-,-,␢,␢)'
				outcomes = outcomes[1:-1].split(',')  # remove parens and get list of faces,
				result = Random.roll(number, len(outcomes))  # then roll for an index in that list
				result = [outcomes[int(r)-1] for r in result.split()]  # that we can then lookup.
				try:  # we don't actually know if the faces are numeric...
					subtotal = sum(map(int, result))
					if plus_minus_dice == '-':
						total -= subtotal
					else:
						total += subtotal
				except (ValueError, TypeError):  # in case faces were not all numeric (for this group or a previous group),
					total = None  # there is no total.
				result = ' '.join(result)  # finally, repack the result
				roll_results.append( (f'{number}d{faces}', f'{result}\n') )  # and append it to output
			elif faces:  # standard dice
				number = int(number) if number else 1  # number may be omitted for single die.
				result = Random.roll(number, faces)
				try:  # this may be part of a larger group of rolls...
					subtotal = sum(map(int, result.split()))
					if plus_minus_dice == '-':
						total -= subtotal  # this may fail because total is None
					else:
						total += subtotal  # this may fail because total is None
				except (ValueError, TypeError):
					total = None
				roll_results.append( (f'{number}d{faces}', result) )  # and append it to output
			elif modifier:  # no dice
				try:  # this may be part of a larger group of rolls...
					if plus_minus_mod == '-':
						total -= int(modifier)  # this may fail because total is None
					elif plus_minus_mod == '+':
						total += int(modifier)  # this may fail because total is None
				except (ValueError, TypeError):
					total = None

		raw_result_list = []
		for raw_roll in roll_results:
			raw_result_list.append(f'=== {raw_roll[0]} ===\n**{raw_roll[1]}**')
		raw = ''.join(raw_result_list)

		output = f"*You rolled {roll} and got **{total}**.*\n\nRaw result:\n{raw}"
		if total is None:
			output = f"*You rolled {roll} and got:*\n\n{raw}"

		return output

	@staticmethod
	def rolls_results(rolls: str) -> str:
		"""Returns a formatted result of dice outcomes"""
		results = [Random.roll_result(roll.strip()) for roll in rolls.split(';')]
		return '\n'.join(results)

	@command(
		name='roll',
		aliases=['dice']
	)
	async def roll_prefix(self, ctx: Context, rolls: str = '1d6'):
		"""Roll some dice (default 1d6)"""
		return await ctx.send(Random.rolls_results(rolls))

	@slash_command(
		name='roll',
		guild_ids=[GUILD]
	)
	async def roll_slash(self,
		ctx: discord.ApplicationContext,
		faces: Option(int, "How many faces are on each die?", min_value=1, max_value=1_000_000_000, default=6),
		number: Option(int, "How many dies should be rolled?", min_value=1, max_value=10_000, default=1),
		modifier: Option(int, "Add this number to the final result", default=0),
	):
		"""Roll some standard dice with an optional modifier (default 1d6 with +0)"""
		if modifier and modifier > 0:
			modifier = f'+{modifier}'
		return await ctx.respond(
			Random.rolls_results(f'{number}d{faces}{modifier}')
		)

# ==============================================================================

	@staticmethod
	def flip(number: int = 1) -> str:
		"""Returns a raw list of Heads or Tails outcomes"""
		result: str = httpx.get(
			f'https://www.random.org/integers/?num={number}&min=1&max=2&col=1&base=10&format=plain&rnd=new'
			).text
		result = result.replace('1', 'Heads')
		result = result.replace('2', 'Tails')
		return result
	
	@staticmethod
	def flip_results(number: int = 1) -> str:
		"""Returns a formatted result of Heads and Tails counts"""
		result: str = Random.flip(number)
		if len(result) == 6: # \n counts as a character and we don't strip it
			return (
				f"*You flipped {number} coins and got:*\n"
				f"**{result}**"
			)
		c = Counter(result.split('\n'))
		return (
			f"*You flipped {number} coins and got:*\n"
			f"**{c['Heads']}** Heads\n"
			f"**{c['Tails']}** Tails"
		)

	@command(
		name='flip',
		aliases=['coin']
	)
	async def flip_prefix(self, ctx: Context, number: str = 1):
		"""Flip a coin (default 1)"""
		return await ctx.send(self.flip_results(number))

	@slash_command(
		name='flip',
		guild_ids=[GUILD]
		)
	async def flip_slash(self,
		ctx: discord.ApplicationContext,
		number: Option(int, "How many coins should be flipped?", min_value=1, max_value=10_000, default=1)
	):
		"""Flip a coin (default 1)"""
		return await ctx.respond(self.flip_results(number))