character.py

from typing import List
from dataclasses import dataclass
import random

@dataclass
class DamageRoll:
    num_dice: int
    dice_sides: int
    modifier: int = 0
    
    def roll(self, is_critical: bool = False) -> tuple[int, list[int]]:
        # On a critical hit, roll dice twice but only add modifier once
        num_rolls = self.num_dice * (2 if is_critical else 1)
        dice_results = [random.randint(1, self.dice_sides) for _ in range(num_rolls)]
        return sum(dice_results) + self.modifier, dice_results

class DamageType:
    SLASHING = "slashing"
    PIERCING = "piercing"
    BLUDGEONING = "bludgeoning"
    FIRE = "fire"
    COLD = "cold"
    LIGHTNING = "lightning"
    THUNDER = "thunder"
    POISON = "poison"

class Effect:
    def __init__(self, name: str, duration: int = None):
        self.name = name
        self.duration = duration
        
    def apply_to_damage(self, damage: int, weapon: 'Weapon', is_critical: bool = False) -> tuple[int, str, list[int]]:
        return damage, None, []
        
    def apply_to_attack(self, bonus: int, weapon: 'Weapon') -> int:
        return bonus

class Rage(Effect):
    def __init__(self):
        super().__init__("Rage", duration=10)  # 10 rounds
        self.rage_damage = 2  # This could scale with barbarian level
        
    def apply_to_damage(self, damage: int, weapon: 'Weapon', is_critical: bool = False) -> tuple[int, str, list[int]]:
        if weapon.is_strength_based:
            return damage + self.rage_damage, weapon.damage_type, []
        return damage, weapon.damage_type, []
        
    def apply_to_attack(self, bonus: int, weapon: 'Weapon') -> int:
        if weapon.is_strength_based:
            return bonus + 2  # Advantage handled separately
        return bonus

    def start_rage(self):
        if not self.is_raging():
            self.add_effect(Rage())
            self.has_advantage_on_strength_checks = True
            self.has_advantage_on_strength_saves = True
            
    def end_rage(self):
        self.remove_effect("Rage")
        self.has_advantage_on_strength_checks = False
        self.has_advantage_on_strength_saves = False

    def get_ability_check_status(self, ability: str = None) -> str:
        """Determine if ability checks should have advantage, disadvantage, or neither"""
        if ability and ability.lower() == "strength" and self.is_raging():
            if self.has_disadvantage_on_ability_checks:
                return 'n'  # Advantage and disadvantage cancel out
            return 'a'
            
        if self.has_advantage_on_ability_checks and self.has_disadvantage_on_ability_checks:
            return 'n'  # They cancel out
        elif self.has_advantage_on_ability_checks:
            return 'a'
        elif self.has_disadvantage_on_ability_checks:
            return 'd'
        return 'n'

    def get_saving_throw_status(self, ability: str) -> tuple[str, bool]:
        """Determine saving throw status and if it should auto-fail
        Returns (advantage_status, auto_fail)"""
        auto_fail = False
        advantage_status = 'n'

        if ability.lower() == "strength":
            if self.auto_fail_strength_saves:
                auto_fail = True
            elif self.has_advantage_on_strength_saves:
                advantage_status = 'a'

        elif ability.lower() == "dexterity":
            if self.auto_fail_dexterity_saves:
                auto_fail = True
            elif self.has_advantage_on_dex_saves and self.has_disadvantage_on_dex_saves:
                advantage_status = 'n'
            elif self.has_advantage_on_dex_saves:
                advantage_status = 'a'
            elif self.has_disadvantage_on_dex_saves:
                advantage_status = 'd'

        return advantage_status, auto_fail

class ElementalWeaponEffect(Effect):
    def __init__(self, element_type: str):
        super().__init__(f"Elemental Weapon: {element_type}", duration=10)
        self.element_type = element_type
        self.elemental_damage = DamageRoll(1, 6)  # 1d6 elemental damage
        
    def apply_to_damage(self, damage: int, weapon: 'Weapon', is_critical: bool = False) -> tuple[int, str, list[int]]:
        # Pass the is_critical flag to roll() so dice get doubled on crit
        damage, dice_rolls = self.elemental_damage.roll(is_critical)
        return damage, self.element_type, dice_rolls

class Weapon:
    def __init__(self, name: str, damage_roll: DamageRoll, attack_bonus: int, 
                 damage_type: str = DamageType.SLASHING, is_strength_based: bool = True):
        self.name = name
        self.damage_roll = damage_roll
        self.attack_bonus = attack_bonus
        self.damage_type = damage_type
        self.is_strength_based = is_strength_based
        self.effects: List[Effect] = []
        
    def add_effect(self, effect: Effect):
        self.effects.append(effect)
        
    def remove_effect(self, effect_name: str):
        self.effects = [e for e in self.effects if e.name != effect_name]
        
    def roll_attack(self, advantage_status: str = 'n', character: 'Character' = None) -> tuple[int, list[int], bool]:
        roll1 = random.randint(1, 20)
        roll2 = random.randint(1, 20) if advantage_status in ['a', 'd'] else None
        
        # Calculate base attack bonus including effects
        bonus = self.attack_bonus
        if character:
            for effect in character.active_effects + self.effects:
                bonus = effect.apply_to_attack(bonus, self)
        
        # Determine final roll and check for critical
        if advantage_status == 'a':
            base_roll = max(roll1, roll2)
            print(f"Attack rolls: {roll1} and {roll2}, using {base_roll} (advantage)")
        elif advantage_status == 'd':
            base_roll = min(roll1, roll2)
            print(f"Attack rolls: {roll1} and {roll2}, using {base_roll} (disadvantage)")
        else:
            base_roll = roll1
            print(f"Attack roll: {base_roll}")
            
        print(f"Adding {bonus} to hit")
            
        is_critical = base_roll == 20
        attack_roll = base_roll + bonus
            
        return attack_roll, [roll1, roll2] if roll2 else [roll1], is_critical
        
    def roll_damage(self, character: 'Character' = None, is_critical: bool = False) -> tuple[int, dict[str, tuple[int, list[int]]]]:
        # Initialize damage breakdown with both total and individual rolls
        base_damage, base_rolls = self.damage_roll.roll(is_critical)
        modifier = self.damage_roll.modifier
        print(f"\n{self.name} damage:")
        print(f"Rolled {base_rolls} + {modifier} {self.damage_type}")
        
        # Calculate total base damage (sum of rolls + modifier)
        total_base = sum(base_rolls) + modifier
        damage_breakdown = {self.damage_type: (total_base, base_rolls)}
        
        # Apply effects from both character and weapon
        if character:
            all_effects = character.active_effects + self.effects
            for effect in all_effects:
                additional_damage, damage_type, dice_rolls = effect.apply_to_damage(0, self, is_critical)
                if additional_damage > 0 and damage_type:
                    if dice_rolls:  # Only print if there were actual dice rolled
                        print(f"Effect {effect.name}: Rolled {dice_rolls} {damage_type}")
                    if damage_type == self.damage_type:
                        # Add to existing damage of same type
                        current_damage, current_rolls = damage_breakdown[damage_type]
                        damage_breakdown[damage_type] = (current_damage + additional_damage, current_rolls)
                    else:
                        # New damage type
                        damage_breakdown[damage_type] = (additional_damage, dice_rolls)
        
        # Print damage totals by type
        print("\nDamage Totals:")
        for damage_type, (damage, _) in damage_breakdown.items():
            print(f"{damage_type}: {damage}")
        print(f"Total damage: {sum(damage[0] for damage in damage_breakdown.values())}")
                    
        return sum(damage[0] for damage in damage_breakdown.values()), damage_breakdown

class Character:
    def __init__(self, name: str):
        self.name = name
        self.weapons: List[Weapon] = []

        self.active_effects: List[Effect] = []
        self.attacks_per_turn = 2
        self.raging_bonus = 2
        # Base stats
        self.proficiency_bonus = 3  # Assumes level 5-8
        
        # Ability score modifiers
        self.strength_mod = 4       # 19 Strength
        self.dexterity_mod = 2      # 14 Dexterity
        self.constitution_mod = 3    # 16 Constitution
        self.intelligence_mod = -1    # 8 Intelligence
        self.wisdom_mod = 0         # 10 Wisdom
        self.charisma_mod = 0      # 10 Charisma
        
        # Saving throw proficiencies (True if proficient)
        self.saving_throws = {
            "strength": True,      # Barbarian is proficient in Str saves
            "dexterity": False,
            "constitution": True,  # Barbarian is proficient in Con saves
            "intelligence": False,
            "wisdom": False,
            "charisma": False
        }
        
        # Saving throw advantages
        self.saving_throw_advantages = {
            "strength": False,
            "dexterity": True,  # Your character has advantage on Dex saves
            "constitution": False,
            "intelligence": False,
            "wisdom": False,
            "charisma": False
        }
        
        # Skill proficiencies (True if proficient)
        self.skill_proficiencies = {
            # Strength skills
            "athletics": True,
            
            # Dexterity skills
            "acrobatics": False,
            "sleight_of_hand": False,
            "stealth": True,
            
            # Intelligence skills
            "arcana": False,
            "history": False,
            "investigation": False,
            "nature": False,
            "religion": False,
            
            # Wisdom skills
            "animal_handling": False,
            "insight": False,
            "medicine": False,
            "perception": False,
            "survival": False,
            
            # Charisma skills
            "deception": True,
            "intimidation": True,
            "performance": False,
            "persuasion": False
        }
        
        # Map skills to their ability scores
        self.skill_abilities = {
            # Strength skills
            "athletics": "strength",
            
            # Dexterity skills
            "acrobatics": "dexterity",
            "sleight_of_hand": "dexterity",
            "stealth": "dexterity",
            
            # Intelligence skills
            "arcana": "intelligence",
            "history": "intelligence",
            "investigation": "intelligence",
            "nature": "intelligence",
            "religion": "intelligence",
            
            # Wisdom skills
            "animal_handling": "wisdom",
            "insight": "wisdom",
            "medicine": "wisdom",
            "perception": "wisdom",
            "survival": "wisdom",
            
            # Charisma skills
            "deception": "charisma",
            "intimidation": "charisma",
            "performance": "charisma",
            "persuasion": "charisma"
        }
        
        # Status effect tracking
        self.status_effects = set()
        self.has_advantage_on_attacks = False
        self.has_disadvantage_on_attacks = False
        self.has_advantage_on_ability_checks = False
        self.has_disadvantage_on_ability_checks = False
        self.auto_fail_strength_saves = False
        self.auto_fail_dexterity_saves = False
        self.has_advantage_on_dex_saves = False
        self.has_disadvantage_on_dex_saves = False
        self.ac_bonus = 0
        self.speed_multiplier = 1
    
    def get_saving_throw_bonus(self, ability: str) -> int:
        """Calculate saving throw bonus for a given ability"""
        ability_map = {
            "strength": self.strength_mod,
            "dexterity": self.dexterity_mod,
            "constitution": self.constitution_mod,
            "intelligence": self.intelligence_mod,
            "wisdom": self.wisdom_mod,
            "charisma": self.charisma_mod
        }
        
        base_bonus = ability_map.get(ability.lower(), 0)
        if self.saving_throws.get(ability.lower(), False):
            base_bonus += self.proficiency_bonus
            
        return base_bonus
    
    def roll_saving_throw(self, ability: str) -> tuple[int, int, bool]:
        """Roll a saving throw for the given ability
        Returns (total, roll, proficient)"""
        roll = random.randint(1, 20)
        bonus = self.get_saving_throw_bonus(ability)
        is_proficient = self.saving_throws.get(ability.lower(), False)
        return roll + bonus, roll, is_proficient
    
    def add_weapon(self, weapon: Weapon):
        self.weapons.append(weapon)
        
    def get_weapon(self, weapon_name: str) -> Weapon:
        for weapon in self.weapons:
            if weapon.name.lower() == weapon_name.lower():
                return weapon
        raise ValueError(f"No weapon named {weapon_name} found")
        
    def add_effect(self, effect: Effect):
        self.active_effects.append(effect)
        
    def remove_effect(self, effect_name: str):
        self.active_effects = [e for e in self.active_effects if e.name != effect_name]
        
    def is_raging(self) -> bool:
        return any(isinstance(effect, Rage) for effect in self.active_effects)
    
    def start_rage(self):
        if not self.is_raging():
            self.add_effect(Rage())
            
    def end_rage(self):
        self.remove_effect("Rage")
    
    def add_elemental_effect(self, weapon_name: str, element_type: str):
        weapon = self.get_weapon(weapon_name)
        # Remove any existing elemental effects
        weapon.effects = [e for e in weapon.effects if not isinstance(e, ElementalWeaponEffect)]
        # Add new elemental effect
        weapon.add_effect(ElementalWeaponEffect(element_type))
        print(f"Added {element_type} damage to {weapon_name}")

    def add_rage_elemental_effect(self, weapon_name: str, element_type: str):
        """Add elemental damage to a weapon while raging. Only one weapon can have elemental damage."""
        if not self.is_raging():
            print("Must be raging to add elemental damage!")
            return
            
        # Remove any existing elemental effects from all weapons
        for weapon in self.weapons:
            weapon.effects = [e for e in weapon.effects if not isinstance(e, ElementalWeaponEffect)]
            
        # Add new elemental effect to specified weapon
        weapon = self.get_weapon(weapon_name)
        weapon.add_effect(ElementalWeaponEffect(element_type))
        print(f"Added {element_type} damage to {weapon_name}")

    def handle_rage_options(self):
        """Handle rage-related options including elemental damage"""
        rage_status = input("Are you raging? (y/n): ").lower()
        if rage_status == 'y':
            self.start_rage()
            # Ask about elemental damage
            elemental = input("Add elemental damage? (y/n): ").lower()
            if elemental == 'y':
                print("\nAvailable elements:")
                print("1. Fire")
                print("2. Cold")
                print("3. Lightning")
                print("4. Thunder")
                element_choice = input("Choose element (1-4): ")
                element_map = {
                    "1": DamageType.FIRE,
                    "2": DamageType.COLD,
                    "3": DamageType.LIGHTNING,
                    "4": DamageType.THUNDER
                }
                if element_choice in element_map:
                    # Get the active weapon (most recently added)
                    if self.weapons:
                        self.add_rage_elemental_effect(self.weapons[-1].name, element_map[element_choice])

    def get_skill_bonus(self, skill: str) -> int:
        """Calculate bonus for a skill check"""
        if skill not in self.skill_abilities:
            raise ValueError(f"Unknown skill: {skill}")
            
        ability = self.skill_abilities[skill]
        ability_mod = getattr(self, f"{ability}_mod")
        
        if self.skill_proficiencies.get(skill, False):
            return ability_mod + self.proficiency_bonus
        return ability_mod
    
    def roll_skill_check(self, skill: str, advantage_status: str = 'n') -> tuple[int, int, bool]:
        """Roll a skill check
        Returns (total, roll, proficient)"""
        roll1 = random.randint(1, 20)
        roll2 = None
        
        if advantage_status in ['a', 'd']:
            roll2 = random.randint(1, 20)
            base_roll = max(roll1, roll2) if advantage_status == 'a' else min(roll1, roll2)
        else:
            base_roll = roll1
            
        bonus = self.get_skill_bonus(skill)
        is_proficient = self.skill_proficiencies.get(skill, False)
        
        return base_roll + bonus, base_roll, is_proficient 

    def add_status(self, status: str) -> None:
        """Add a status effect and its corresponding modifiers"""
        self.status_effects.add(status)
        
        if status == "blinded":
            self.has_disadvantage_on_attacks = True
            
        elif status == "invisible":
            self.has_advantage_on_attacks = True
            
        elif status == "paralyzed":
            self.auto_fail_strength_saves = True
            self.auto_fail_dexterity_saves = True
            
        elif status == "petrified":
            self.auto_fail_strength_saves = True
            self.auto_fail_dexterity_saves = True
            
        elif status == "poisoned":
            self.has_disadvantage_on_attacks = True
            self.has_disadvantage_on_ability_checks = True
            
        elif status == "prone":
            self.has_disadvantage_on_attacks = True
            
        elif status == "restrained":
            self.has_disadvantage_on_attacks = True
            self.has_disadvantage_on_dex_saves = True
            
        elif status == "stunned":
            self.auto_fail_strength_saves = True
            self.auto_fail_dexterity_saves = True
            
        elif status == "unconscious":
            self.auto_fail_strength_saves = True
            self.auto_fail_dexterity_saves = True
            
        elif status == "hasted":
            self.ac_bonus += 2
            self.has_advantage_on_dex_saves = True
            self.speed_multiplier = 2
            
        elif status == "lethargic":
            self.cannot_take_actions = True
            self.cannot_move = True

    def remove_status(self, status: str) -> None:
        """Remove a status effect and its corresponding modifiers"""
        if status in self.status_effects:
            self.status_effects.remove(status)
            
            if status == "blinded":
                self.has_disadvantage_on_attacks = False
                
            elif status == "invisible":
                self.has_advantage_on_attacks = False
                
            elif status == "paralyzed":
                self.auto_fail_strength_saves = False
                self.auto_fail_dexterity_saves = False
                
            elif status == "petrified":
                self.auto_fail_strength_saves = False
                self.auto_fail_dexterity_saves = False
                
            elif status == "poisoned":
                self.has_disadvantage_on_attacks = False
                self.has_disadvantage_on_ability_checks = False
                
            elif status == "prone":
                self.has_disadvantage_on_attacks = False
                
            elif status == "restrained":
                self.has_disadvantage_on_attacks = False
                self.has_disadvantage_on_dex_saves = False
                
            elif status == "stunned":
                self.auto_fail_strength_saves = False
                self.auto_fail_dexterity_saves = False
                
            elif status == "unconscious":
                self.auto_fail_strength_saves = False
                self.auto_fail_dexterity_saves = False
                
            elif status == "hasted":
                self.ac_bonus -= 2
                self.has_advantage_on_dex_saves = False
                self.speed_multiplier = 1
                
            elif status == "lethargic":
                self.cannot_take_actions = False
                self.cannot_move = False


    def get_attack_roll_status(self) -> str:
        """Determine if attack rolls should have advantage, disadvantage, or neither"""
        if self.has_advantage_on_attacks and self.has_disadvantage_on_attacks:
            return 'n'  # They cancel out
        elif self.has_advantage_on_attacks:
            return 'a'
        elif self.has_disadvantage_on_attacks:
            return 'd'
        return 'n'

    def get_ability_check_status(self, ability: str = None) -> str:
        """Determine if ability checks should have advantage, disadvantage, or neither"""
        if ability and ability.lower() == "strength" and self.is_raging():
            if self.has_disadvantage_on_ability_checks:
                return 'n'  # Advantage and disadvantage cancel out
            return 'a'
            
        if self.has_advantage_on_ability_checks and self.has_disadvantage_on_ability_checks:
            return 'n'  # They cancel out
        elif self.has_advantage_on_ability_checks:
            return 'a'
        elif self.has_disadvantage_on_ability_checks:
            return 'd'
        return 'n'

    def get_saving_throw_status(self, ability: str) -> tuple[str, bool]:
        """Determine saving throw status and if it should auto-fail
        Returns (advantage_status, auto_fail)"""
        auto_fail = False
        advantage_status = 'n'

        if ability.lower() == "dexterity":
            if self.auto_fail_dexterity_saves:
                auto_fail = True
            elif self.has_advantage_on_dex_saves and self.has_disadvantage_on_dex_saves:
                advantage_status = 'n'
            elif self.has_advantage_on_dex_saves:
                advantage_status = 'a'
            elif self.has_disadvantage_on_dex_saves:
                advantage_status = 'd'

        elif ability.lower() == "strength" and self.auto_fail_strength_saves:
            auto_fail = True

        return advantage_status, auto_fail

    def get_current_ac(self) -> int:
        """Get current AC including status effect modifiers"""
        base_ac = 10 + self.dexterity_mod  # Modify this based on your AC calculation
        return base_ac + self.ac_bonus