import os
import re
import random
from typing import Dict, List, Tuple
import discord
from discord import app_commands
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
TOKEN = os.getenv('DISCORD_TOKEN')

# Initialize Discord client
class DiceBot(discord.Client):
    def __init__(self):
        super().__init__(intents=discord.Intents.default())
        self.tree = app_commands.CommandTree(self)

    async def setup_hook(self):
        await self.tree.sync()

client = DiceBot()

# Genesys dice definitions
GENESYS_DICE = {
    'g': [  # Green (Ability)
        ('n', ''),
        ('s', ''),
        ('s', ''),
        ('ss', ''),
        ('a', ''),
        ('a', ''),
        ('s,a', ''),
        ('aa', ''),
    ],
    'y': [  # Yellow (Proficiency)
        ('n', ''),
        ('s', ''),
        ('s', ''),
        ('ss', ''),
        ('ss', ''),
        ('a', ''),
        ('s,a', ''),
        ('s,a', ''),
        ('s,a', ''),
        ('aa', ''),
        ('aa', ''),
        ('t', ''),  # Triumph counts as success
    ],
    'p': [  # Purple (Difficulty)
        ('n', ''),
        ('f', ''),
        ('f,d', ''),
        ('d', ''),
        ('d', ''),
        ('d', ''),
        ('dd', ''),
        ('f', ''),
    ],
    'r': [  # Red (Challenge)
        ('n', ''),
        ('f', ''),
        ('f', ''),
        ('f,d', ''),
        ('f,d', ''),
        ('d', ''),
        ('d', ''),
        ('dd', ''),
        ('dd', ''),
        ('d,f', ''),
        ('d,f', ''),
        ('x', ''),  # Despair counts as failure
    ],
    'b': [  # Blue (Boost)
        ('n', ''),
        ('n', ''),
        ('s', ''),
        ('s,a', ''),
        ('aa', ''),
        ('a', ''),
    ],
    'k': [  # Black (Setback)
        ('', ''),
        ('', ''),
        ('f', ''),
        ('f', ''),
        ('d', ''),
        ('d', ''),
    ]
}

def parse_dice_notation(notation: str) -> List[Tuple[int, str]]:
    """Parse dice notation into a list of (count, die_type) tuples."""
    parts = notation.lower().replace(' ', '').split('+')
    dice = []
    
    for part in parts:
        match = re.match(r'(\d+)d([0-9gpbkry]+)', part)
        if match:
            count = int(match.group(1))
            die_type = match.group(2)
            dice.append((count, die_type))
    
    return dice

def roll_traditional_dice(count: int, sides: int) -> List[int]:
    """Roll traditional dice."""
    return [random.randint(1, sides) for _ in range(count)]

def roll_popping_dice(count: int, sides: int) -> List[List[int]]:
    """Roll dice that spawn a d4 on rolling 1.
    Returns a list of lists, where each inner list represents the chain of rolls from one initial die."""
    results = []
    for _ in range(count):
        chain = []
        roll = random.randint(1, sides)
        chain.append(roll)
        while roll == 1:  # Keep rolling d4s as long as we get 1s
            roll = random.randint(1, 4)
            chain.append(roll)
        results.append(chain)
    return results

def roll_genesys_die(die_type: str) -> Tuple[str, str]:
    """Roll a single Genesys die and return its result."""
    return random.choice(GENESYS_DICE[die_type])

def calculate_genesys_results(results: List[Tuple[str, str]]) -> Dict[str, int]:
    """Calculate net results from Genesys dice rolls."""
    success = 0
    advantage = 0
    triumph = 0
    despair = 0
    blank = 0
    
    for result, _ in results:
        for symbol in result.split(','):
            if symbol == 's':
                success += 1
            elif symbol == 'f':
                success -= 1
            elif symbol == 'a':
                advantage += 1
            elif symbol == 'd':
                advantage -= 1
            elif symbol == 't':
                triumph += 1
                success += 1
            elif symbol == 'x':
                despair += 1
                success -= 1
            elif symbol == 'n':
                blank += 1
    
    return {
        'success': success,
        'advantage': advantage,
        'triumph': triumph,
        'despair': despair,
        'blank': blank
    }

@client.tree.command(name="roll", description="Roll traditional dice (e.g., 2d6 + 1d8)")
async def roll(interaction: discord.Interaction, dice: str):
    try:
        dice_sets = parse_dice_notation(dice)
        if not dice_sets:
            await interaction.response.send_message("Invalid dice notation. Use format like '2d6 + 1d8'")
            return

        results = []
        total = 0
        
        for count, die in dice_sets:
            if not die.isdigit():
                await interaction.response.send_message(f"Invalid die type: {die} did you include +? Use format like '2d6 + 1d8'")
                return
                
            sides = int(die)
            rolls = roll_traditional_dice(count, sides)
            results.append(f"{count}d{sides}: {rolls} = {sum(rolls)}")
            total += sum(rolls)

        response = "\n".join(results)
        if len(dice_sets) > 1:
            response += f"\nTotal: {total}"
        
        await interaction.response.send_message(response)
    except Exception as e:
        await interaction.response.send_message(f"Error: {str(e)}")

@client.tree.command(name="genesys", description="Roll Genesys/Star Wars dice (e.g., 2dg + 1dy)")
async def genesys(interaction: discord.Interaction, dice: str):
    try:
        dice_sets = parse_dice_notation(dice)
        if not dice_sets:
            await interaction.response.send_message("Invalid dice notation. Use format like '2dg + 1dy'")
            return

        all_results = []
        for count, die_type in dice_sets:
            if die_type not in GENESYS_DICE:
                await interaction.response.send_message(f"Invalid die type: {die_type}, did you include +? Use format like '2dg + 1dy'")
                return
            
            for _ in range(count):
                all_results.append(roll_genesys_die(die_type))

        net_results = calculate_genesys_results(all_results)
        
        response = []
        if net_results['success'] > 0:
            response.append(f"{net_results['success']} Success{'es' if net_results['success'] > 1 else ''}")
        elif net_results['success'] < 0:
            response.append(f"{abs(net_results['success'])} Failure{'s' if abs(net_results['success']) > 1 else ''}")
            
        if net_results['advantage'] > 0:
            response.append(f"{net_results['advantage']} Advantage{'s' if net_results['advantage'] > 1 else ''}")
        elif net_results['advantage'] < 0:
            response.append(f"{abs(net_results['advantage'])} Disadvantage{'s' if abs(net_results['advantage']) > 1 else ''}")
            
        if net_results['triumph'] > 0:
            response.append(f"{net_results['triumph']} Triumph{'s' if net_results['triumph'] > 1 else ''}")
        if net_results['despair'] > 0:
            response.append(f"{net_results['despair']} Despair{'s' if net_results['despair'] > 1 else ''}")
        if net_results['blank'] > 0:
            response.append(f"{net_results['blank']} Blank")

        await interaction.response.send_message(" | ".join(response) if response else "Neutral Result, No Blanks")
    except Exception as e:
        await interaction.response.send_message(f"Error: {str(e)}")

@client.tree.command(name="pop", description="Roll dice that spawn 1d4 on rolling 1 (e.g., 2d6 + 1d8)")
async def pop(interaction: discord.Interaction, dice: str):
    try:
        dice_sets = parse_dice_notation(dice)
        if not dice_sets:
            await interaction.response.send_message("Invalid dice notation. Use format like '2d6 + 1d8'")
            return

        results = []
        total = 0
        
        for count, die in dice_sets:
            if not die.isdigit():
                await interaction.response.send_message(f"Invalid die type: {die} did you include +? Use format like '2d6 + 1d8'")
                return
                
            sides = int(die)
            rolls = roll_popping_dice(count, sides)
            
            # Format the output for this dice set
            roll_strs = []
            set_total = 0
            for chain in rolls:
                chain_str = str(chain[0])
                if len(chain) > 1:
                    chain_str += f" → {' → '.join(map(str, chain[1:]))}"
                roll_strs.append(chain_str)
                set_total += sum(chain)
            
            results.append(f"{count}d{sides}: [{', '.join(roll_strs)}] = {set_total}")
            total += set_total

        response = "\n".join(results)
        if len(dice_sets) > 1:
            response += f"\nTotal: {total}"
        
        await interaction.response.send_message(response)
    except Exception as e:
        await interaction.response.send_message(f"Error: {str(e)}")

@client.event
async def on_ready():
    print(f'{client.user} has connected to Discord!')

if __name__ == "__main__":
    client.run(TOKEN)