from .base_agent import BaseAgent
from ..aws_connector import AWSConnector
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta

class CostManagementAgent(BaseAgent):
    """
    An agent specialized in handling AWS Cost Management and Optimization tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Cost Management.
        """
        if command == 'list_budgets':
            return self._list_budgets()
        elif command == 'get_monthly_cost':
            return self._get_monthly_cost()
        elif command == 'get_optimization_suggestions':
            return self._get_optimization_suggestions()
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by CostManagementAgent.")

    def _list_budgets(self):
        """Lists all AWS Budgets for the account."""
        # ... (code is unchanged)

    def _get_monthly_cost(self):
        """Gets the total cost for the last full month."""
        # ... (code is unchanged)

    def _get_optimization_suggestions(self):
        """Finds potential cost-saving opportunities in the AWS account."""
        print("CostManagementAgent: Analyzing account for cost optimization suggestions...")
        suggestions = []
        try:
            # Get all available regions for EC2
            ec2_client = AWSConnector.get_client('ec2')
            regions = [region['RegionName'] for region in ec2_client.describe_regions()['Regions']]

            for region in regions:
                print(f"Analyzing region: {region}...")
                region_client = AWSConnector.get_client('ec2', region_name=region)

                # Find unattached EBS volumes
                volumes = region_client.describe_volumes(Filters=[{'Name': 'status', 'Values': ['available']}]).get('Volumes', [])
                for volume in volumes:
                    suggestion = f"Unattached EBS Volume: {volume['VolumeId']} in region {region} is 'available' and could be deleted to save costs."
                    suggestions.append(suggestion)

                # Find unassociated Elastic IPs
                addresses = region_client.describe_addresses().get('Addresses', [])
                for address in addresses:
                    if 'AssociationId' not in address:
                        suggestion = f"Unassociated Elastic IP: {address['PublicIp']} in region {region} is not associated with any resource and is incurring costs."
                        suggestions.append(suggestion)
            
            if not suggestions:
                return {"status": "success", "message": "No immediate cost optimization suggestions found for unattached EBS volumes or unassociated Elastic IPs."}

            return {"status": "success", "suggestions": suggestions}
        except Exception as e:
            print(f"Error getting optimization suggestions: {e}")
            return {"status": "error", "message": str(e)}