combat_manager.py

from character import Character, DamageType
from initiative import roll_initiative
from saving_throw import roll_saving_throw
from skill_check import roll_skill_check
from weapons.registry import WeaponRegistry
from weapons.melee.improvised_weapon import ImprovisedWeapon
from typing import List, Optional
from dataclasses import dataclass
from datetime import datetime

@dataclass
class StatusEffect:
    name: str
    duration: Optional[int]  # Number of turns, None for indefinite
    start_turn: int  # Turn number when effect was applied

class CombatManager:
    def __init__(self, character: Character):
        self.character = character
        self.is_turn_active = False
        self.attacks_remaining = 0
        self.bonus_action_available = True
        self.character.bonus_action_available = True
        self.inspiration_available = True
        self.available_elements = {
            "1": DamageType.FIRE,
            "2": DamageType.COLD,
            "3": DamageType.LIGHTNING,
            "4": DamageType.THUNDER
        }
        
        # Add default weapons
        default_weapons = ["tavern_brawler", "poison_dagger", "javelin_of_lightning", "vicious_scimitar"]
        for weapon_id in default_weapons:
            try:
                weapon = WeaponRegistry.get_weapon(weapon_id)
                if weapon:
                    self.character.add_weapon(weapon)
            except ValueError as e:
                print(f"Warning: Could not add default weapon {weapon_id}: {e}")

        self.status_effects = {}  # Dictionary of StatusEffect objects
        self.status_definitions = {
            "blinded": "Can't see, auto-fails sight checks. Attack rolls against have advantage, attacks have disadvantage.",
            "charmed": "Can't attack charmer or target them with harmful effects. Charmer has advantage on social checks.",
            "deafened": "Can't hear, auto-fails hearing checks.",
            "frightened": "Disadvantage on ability checks and attacks while source in sight. Can't move closer to source.",
            "grappled": "Speed becomes 0, no bonus to speed.",
            "incapacitated": "Can't take actions or reactions.",
            "invisible": "Impossible to see without magic/special sense. Attack rolls against have disadvantage, attacks have advantage.",
            "paralyzed": "Incapacitated, can't move/speak. Auto-fails STR/DEX saves. Attacks have advantage, crits within 5ft.",
            "petrified": "Turned to stone, incapacitated. Advantage on attacks against, auto-fails STR/DEX saves. Resistant to all damage.",
            "poisoned": "Disadvantage on attack rolls and ability checks.",
            "prone": "Can only crawl. Disadvantage on attacks. Melee attacks have advantage, ranged disadvantage against.",
            "restrained": "Speed 0, no bonus. Advantage on attacks against, disadvantage on attacks. Disadvantage on DEX saves.",
            "stunned": "Incapacitated, can't move, faltering speech. Auto-fails STR/DEX saves. Advantage on attacks against.",
            "unconscious": "Incapacitated, unaware. Drops items, falls prone. Auto-fails STR/DEX saves. Advantage on attacks, crits within 5ft.",
            "hasted": "Double speed, +2 AC, advantage on DEX saves, extra limited action. Lethargy when ended."
        }
        self.current_turn = 0

    def start_combat(self):
        """Initialize combat by rolling initiative"""
        print("\n=== COMBAT STARTS ===")
        roll_initiative(self.character)
        self.is_turn_active = False

    def start_turn(self):
        """Reset available actions for new turn"""
        print("\n=== YOUR TURN ===")
        self.current_turn += 1
        self.is_turn_active = True
        self.attacks_remaining = self.character.attacks_per_turn
            
        self.bonus_action_available = True
        self.character.bonus_action_available = True
        
        # Check for expired effects
        expired_effects = []
        for status, effect in self.status_effects.items():
            if effect.duration and (self.current_turn - effect.start_turn) >= effect.duration:
                expired_effects.append(status)
                
        # Handle expired effects
        for status in expired_effects:
            if status == "hasted":
                print("\nHaste has expired! You are now lethargic.")
                self.remove_status_effect(status)
                self.add_status_effect("lethargic")
            else:
                print(f"\n{status.capitalize()} effect has expired!")
                self.remove_status_effect(status)
        
        # Handle lethargy
        if "lethargic" in self.status_effects:
            print("\nYou are lethargic from Haste wearing off!")
            print("You can't move or take actions this turn!")
            self.attacks_remaining = 0
            self.bonus_action_available = False
            self.remove_status_effect("lethargic")  # Removes after one turn
            
        self.print_status()

    def print_status(self):
        """Display current combat status"""
        print("\nStatus:")
        print(f"Current Turn: {self.current_turn}")
        print(f"Attacks remaining: {self.attacks_remaining}")
        print(f"Bonus action: {'Available' if self.bonus_action_available else 'Used'}")
        print(f"Inspiration: {'Available' if self.inspiration_available else 'Used'}")
        print(f"Raging: {'Yes' if self.character.is_raging() else 'No'}")
        
        if self.status_effects:
            print("\nActive Status Effects:")
            for status, effect in sorted(self.status_effects.items()):
                turns_remaining = effect.duration - (self.current_turn - effect.start_turn) if effect.duration else "∞"
                print(f"- {status.capitalize()} ({turns_remaining} turns remaining)")
        
        # Show active weapons and their effects
        print("\nWeapons:")
        for weapon in self.character.weapons:
            weapon_desc = weapon.name
            if weapon.elemental_damage:
                weapon_desc += f" (+1d6 {weapon.elemental_damage} damage)"
            effects = [e.name for e in weapon.effects]
            if effects:
                weapon_desc += f" ({', '.join(effects)})"
            print(f"- {weapon_desc}")

    def handle_action(self):
        """Present and handle available actions"""
        while self.is_turn_active:
            print("\nAvailable Actions:")
            options = []
            
            if self.attacks_remaining > 0:
                options.append("1. Attack")
            if self.bonus_action_available:
                options.append("2. Bonus Action")
            options.extend([
                "3. Check Status",
                "4. Make Ability Check",
                "5. Make Saving Throw",
                "6. Manage Status Effects",
                "7. Manage Weapons",
                "8. End Turn"
            ])
            
            for option in options:
                print(option)

            choice = input("\nChoose action: ")

            if choice == "1" and self.attacks_remaining > 0:
                self.handle_attack()
            elif choice == "2" and self.bonus_action_available:
                self.handle_bonus_action()
            elif choice == "3":
                self.print_status()
            elif choice == "4":
                self.handle_ability_check()
            elif choice == "5":
                self.handle_saving_throw()
            elif choice == "6":
                self.manage_status_effects()
            elif choice == "7":
                self.manage_weapons()
            elif choice == "8":
                self.end_turn()
                break
            else:
                print("Invalid choice or action not available")

    def handle_attack(self):
        """Handle attack action"""
        print("\nChoose weapon:")
        for i, weapon in enumerate(self.character.weapons, 1):
            # Create weapon description with elemental effect if present
            weapon_desc = weapon.name
            if weapon.elemental_damage:
                weapon_desc += f" (+1d6 {weapon.elemental_damage} damage)"
            # Add any other effects the weapon might have
            effects = [e.name for e in weapon.effects]
            if effects:
                weapon_desc += f" ({', '.join(effects)})"
            print(f"{i}. {weapon_desc}")
        
        weapon_choice = int(input("Select weapon: ")) - 1
        if 0 <= weapon_choice < len(self.character.weapons):
            weapon = self.character.weapons[weapon_choice]
            
            # Check for special actions
            special_actions = weapon.get_special_actions()
            if special_actions:
                print("\nSpecial Actions Available:")
                for i, action in enumerate(special_actions, 1):
                    print(f"{i}. {action}")
                print(f"{len(special_actions) + 1}. Normal Attack")
                
                choice = input("\nChoose action: ")
                if choice.isdigit() and 1 <= int(choice) <= len(special_actions):
                    weapon.handle_special_action(special_actions[int(choice)-1], self.character)
                    self.attacks_remaining -= 1
                    return
            
            # Normal attack procedure
            # Check character's attack roll status first
            attack_status = self.character.get_attack_roll_status()
            
            if attack_status == 'n':
                print("\nRoll options:")
                print("(a) Advantage")
                print("(d) Disadvantage") 
                print("(n) Normal")
                if self.inspiration_available:
                    print("(i) Use Inspiration (automatic critical hit)")
                advantage = input("Choose roll type: ").lower()
            else:
                print(f"\nAttack roll automatically has {'advantage' if attack_status == 'a' else 'disadvantage'} due to status effects")
                advantage = attack_status
            
            if advantage == 'i' and self.inspiration_available:
                # Automatic critical hit
                print("Using inspiration for automatic critical hit!")
                attack_roll = 20  # Natural 20
                is_critical = True
                self.inspiration_available = False
                print(f"Attack roll: {attack_roll} (Natural 20!)")
                print("CRITICAL HIT!")
                hit = 'y'
            else:
                # Normal attack roll logic
                attack_roll, dice_rolls, is_critical = weapon.roll_attack(advantage, self.character)
                
                if is_critical:
                    print("CRITICAL HIT!")
                    hit = 'y'
                else:
                    hit = input("Does the attack hit? (y/n): ").lower()
            
            if hit == 'y':
                total_damage, damage_breakdown = weapon.roll_damage(self.character, is_critical)
                
                # Print damage breakdown with dice rolls
                for damage_type, (damage, dice_rolls) in damage_breakdown.items():
                    if dice_rolls:
                        print(f"{damage} {damage_type} damage (rolled {dice_rolls})")
                    else:
                        print(f"{damage} {damage_type} damage")
                
                print(f"\nTotal damage: {total_damage}")
                
                # Call handle_hit for any weapon-specific hit effects
                weapon.handle_hit(self.character, is_critical)
                
                if self.character.bonus_action_available == False:
                    self.bonus_action_available = False
            
            self.attacks_remaining -= 1
        else:
            print("Invalid weapon choice")

    def handle_bonus_action(self):
        """Handle bonus action options"""
        print("\nBonus Action Options:")
        if self.character.is_raging() == False:
            print("1. Start Rage")
        if self.character.is_raging():
            print("2. Add/Change Elemental Effect")
        print("3. Make Ability Check")
        print("4. Cancel")
        
        choice = input("Choose option: ")
        
        if choice == "1" and not self.character.is_raging():
            self.character.start_rage()
            self.bonus_action_available = False
            self.character.bonus_action_available = False
            print("You are now raging!")
            
            # Immediately prompt for elemental effect when starting rage
            print("\nWould you like to add an elemental effect to a weapon? (y/n): ")
            if input().lower() == 'y':
                self._handle_elemental_effect()
            
        elif choice == "2" and self.character.is_raging():
            self._handle_elemental_effect()
            
        elif choice == "3":
            self.handle_ability_check()
            self.bonus_action_available = False
            self.character.bonus_action_available = False
            
    
    def _handle_elemental_effect(self):
        """Helper method to handle adding/changing elemental effects"""
        print("\nChoose weapon to enchant:")
        for i, weapon in enumerate(self.character.weapons, 1):
            print(f"{i}. {weapon.name}")
        
        weapon_choice = int(input("Select weapon: ")) - 1
        if 0 <= weapon_choice < len(self.character.weapons):
            print("\nChoose element:")
            for num, element in self.available_elements.items():
                print(f"{num}. {element}")
            
            element_choice = input("Select element: ")
            if element_choice in self.available_elements:
                weapon = self.character.weapons[weapon_choice]
                self.character.add_rage_elemental_effect(
                    weapon.name, 
                    self.available_elements[element_choice]
                )
                self.bonus_action_available = False
                self.character.bonus_action_available = False
            else:
                print("Invalid element choice")
        else:
            print("Invalid weapon choice")

    def handle_ability_check(self):
        """Handle ability check rolls"""
        print("\nChoose skill:")
        skills = list(self.character.skill_abilities.keys())
        for i, skill in enumerate(skills, 1):
            print(f"{i}. {skill.replace('_', ' ').capitalize()}")
        
        skill_choice = int(input("Select skill: ")) - 1
        if 0 <= skill_choice < len(skills):
            chosen_skill = skills[skill_choice]
            
            # Get the ability associated with this skill
            ability = self.character.skill_abilities[chosen_skill]
            
            # Check if status effects determine advantage/disadvantage
            advantage_status = self.character.get_ability_check_status(ability)
            
            if advantage_status != 'n':
                print(f"Rolling with {'advantage' if advantage_status == 'a' else 'disadvantage'} due to status effects")
                advantage = advantage_status
            else:
                advantage = input("Roll with (a)dvantage, (d)isadvantage, or (n)ormal?: ").lower()
                
            roll_skill_check(self.character, chosen_skill, advantage)
        else:
            print("Invalid skill choice")

    def handle_saving_throw(self):
        """Handle saving throw rolls"""
        print("\nChoose ability:")
        abilities = ["strength", "dexterity", "constitution", "intelligence", "wisdom", "charisma"]
        for i, ability in enumerate(abilities, 1):
            print(f"{i}. {ability.capitalize()}")
        
        ability_choice = int(input("Select ability: ")) - 1
        if 0 <= ability_choice < len(abilities):
            chosen_ability = abilities[ability_choice]
            
            # Check if status effects determine advantage/disadvantage
            advantage_status, auto_fail = self.character.get_saving_throw_status(chosen_ability)
            
            if auto_fail:
                print(f"Auto-failing {chosen_ability} save due to status effects")
                return
            elif advantage_status != 'n':
                print(f"Rolling with {'advantage' if advantage_status == 'a' else 'disadvantage'} due to status effects")
                advantage = advantage_status
            else:
                advantage = input("Roll with (a)dvantage, (d)isadvantage, or (n)ormal?: ").lower()
                
            roll_saving_throw(self.character, chosen_ability, advantage)
        else:
            print("Invalid ability choice")

    def manage_status_effects(self):
        """Manage character status effects"""
        while True:
            print("\nStatus Effect Management:")
            print("1. Add Status Effect")
            print("2. Remove Status Effect")
            print("3. List Active Status Effects")
            print("4. View Status Effect Descriptions")
            print("5. Back")

            choice = input("\nChoose option: ")

            if choice == "1":
                print("\nAvailable Status Effects:")
                available_statuses = sorted(self.status_definitions.keys())
                for i, status in enumerate(available_statuses, 1):
                    print(f"{i}. {status.capitalize()}")
                
                status_choice = int(input("\nSelect status effect to add: ")) - 1
                if 0 <= status_choice < len(available_statuses):
                    chosen_status = available_statuses[status_choice]
                    self.add_status_effect(chosen_status)
                else:
                    print("Invalid status choice")
            elif choice == "2":
                self.remove_status_effect()
            elif choice == "3":
                self.list_active_status_effects()
            elif choice == "4":
                self.view_status_descriptions()
            elif choice == "5":
                break
            else:
                print("Invalid choice")

    def add_status_effect(self, status: str):
        """Add a status effect to the character"""
        if status == "hasted":
            duration = 10
        elif status == "lethargic":
            duration = 1
        else:
            duration = None  # Indefinite duration for other effects
            
        self.status_effects[status] = StatusEffect(
            name=status,
            duration=duration,
            start_turn=self.current_turn
        )
        self.character.add_status(status)
        
        if status == "hasted":
            print("\nYou gain the following benefits from Haste:")
            print("- Your speed is doubled")
            print("- You gain a +2 bonus to AC")
            print("- You have advantage on Dexterity saving throws")
            print("- You gain an additional action each turn")
            print("Warning: When Haste ends, you will be lethargic for one turn!")
            self.character.attacks_per_turn += 1
            self.attacks_remaining += 1

    def remove_status_effect(self, status: str):
        """Remove a status effect from the character"""
        if status in self.status_effects:
            del self.status_effects[status]
            self.character.remove_status(status)

    def list_active_status_effects(self):
        """List all active status effects"""
        if not self.status_effects:
            print("\nNo active status effects")
            return

        print("\nActive Status Effects:")
        for status, effect in sorted(self.status_effects.items()):
            turns_remaining = effect.duration - (self.current_turn - effect.start_turn) if effect.duration else "∞"
            print(f"- {status.capitalize()} ({turns_remaining} turns remaining)")

    def view_status_descriptions(self):
        """View descriptions of all status effects"""
        print("\nStatus Effect Descriptions:")
        for status, description in sorted(self.status_definitions.items()):
            print(f"\n{status.capitalize()}:")
            print(description)

    def end_turn(self):
        """End the current turn"""
        self.is_turn_active = False
        print("\n=== END OF TURN ===")

    def manage_weapons(self):
        """Manage available weapons"""
        while True:
            print("\nWeapon Management:")
            print("1. Add Weapon")
            print("2. Remove Weapon")
            print("3. List Current Weapons")
            print("4. List All Available Weapons")
            print("5. Back")

            choice = input("\nChoose option: ")

            if choice == "1":
                self.add_weapon()
            elif choice == "2":
                self.remove_weapon()
            elif choice == "3":
                self.list_current_weapons()
            elif choice == "4":
                self.list_available_weapons()
            elif choice == "5":
                break
            else:
                print("Invalid choice")

    def add_weapon(self):
        """Add a weapon to the character's arsenal"""
        print("\nAvailable Weapons:")
        all_weapons = sorted(WeaponRegistry.get_all_weapon_ids())
        current_weapons = {w.name.split(' (')[0] for w in self.character.weapons}  # Handle descriptors
        
        available_weapons = []
        for i, weapon_id in enumerate(all_weapons, 1):
            # Special handling for improvised weapons - always show as available
            if weapon_id == "improvised_weapon":
                available_weapons.append((i, weapon_id, "Improvised Weapon (create new)"))
                print(f"{i}. Improvised Weapon (create new)")
                continue
                
            # Regular weapons
            weapon = WeaponRegistry.get_weapon(weapon_id)
            if weapon and weapon.name not in current_weapons:
                available_weapons.append((i, weapon_id, weapon.name))
                print(f"{i}. {weapon.name}")
        
        if not available_weapons:
            print("No additional weapons available")
            return

        try:
            choice = int(input("\nSelect weapon to add: ")) - 1
            if 0 <= choice < len(available_weapons):
                _, weapon_id, weapon_name = available_weapons[choice]
                
                # Create the weapon
                if weapon_id == "improvised_weapon":
                    descriptor = input("\nWhat are you using as an improvised weapon? ")
                    weapon = ImprovisedWeapon(descriptor)
                else:
                    weapon = WeaponRegistry.get_weapon(weapon_id)
                
                self.character.add_weapon(weapon)
                print(f"\nAdded {weapon.name}")
            else:
                print("Invalid choice")
        except ValueError:
            print("Please enter a valid number")

    def remove_weapon(self):
        """Remove a weapon from the character's arsenal"""
        if not self.character.weapons:
            print("\nNo weapons to remove")
            return

        print("\nCurrent Weapons:")
        for i, weapon in enumerate(self.character.weapons, 1):
            weapon_desc = weapon.name
            if weapon.elemental_damage:
                weapon_desc += f" (+1d6 {weapon.elemental_damage} damage)"
            print(f"{i}. {weapon_desc}")

        try:
            choice = int(input("\nSelect weapon to remove: ")) - 1
            if 0 <= choice < len(self.character.weapons):
                weapon = self.character.weapons[choice]
                self.character.weapons.remove(weapon)
                print(f"\nRemoved {weapon.name}")
            else:
                print("Invalid choice")
        except ValueError:
            print("Please enter a valid number")

    def list_current_weapons(self):
        """List all weapons currently available to the character"""
        if not self.character.weapons:
            print("\nNo weapons currently available")
            return

        print("\nCurrent Weapons:")
        for weapon in self.character.weapons:
            weapon_desc = weapon.name
            if weapon.elemental_damage:
                weapon_desc += f" (+1d6 {weapon.elemental_damage} damage)"
            effects = [e.name for e in weapon.effects]
            if effects:
                weapon_desc += f" ({', '.join(effects)})"
            print(f"- {weapon_desc}")

    def list_available_weapons(self):
        """List all weapons in the registry"""
        print("\nAll Available Weapons:")
        for weapon_id in sorted(WeaponRegistry.get_all_weapon_ids()):
            weapon = WeaponRegistry.get_weapon(weapon_id)
            if weapon:
                print(f"- {weapon.name}")

def main():
    character = Character("Your Character")
    combat = CombatManager(character)
    
    print("=== COMBAT MANAGER ===")
    combat.start_combat()
    while True:
        print("\nOptions:")
        print("1. Start Your Turn")
        print("2. Make Saving Throw")
        print("3. Make Ability Check")
        print("4. Manage Status Effects")
        print("5. Manage Weapons")
        print("6. End Combat")
        
        choice = input("\nChoose option: ")
        
        if choice == "1":
            combat.start_turn()
            combat.handle_action()
        elif choice == "2":
            combat.handle_saving_throw()
        elif choice == "3":
            combat.handle_ability_check()
        elif choice == "4":
            combat.manage_status_effects()
        elif choice == "5":
            combat.manage_weapons()
        elif choice == "6":
            print("\n=== COMBAT ENDS ===")
            break
        else:
            print("Invalid choice")

if __name__ == "__main__":
    main()