from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class CloudTrailAgent(BaseAgent):
    """
    An agent specialized in handling AWS CloudTrail tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to CloudTrail.
        """
        if command == 'list_trails':
            return self._list_trails(**kwargs)
        elif command == 'describe_trails':
            return self._describe_trails(**kwargs)
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by CloudTrailAgent.")

    def _list_trails(self, region: str):
        """Lists all CloudTrail trails in a specified region."""
        print(f"CloudTrailAgent: Listing trails in region {region}...")
        try:
            cloudtrail_client = AWSConnector.get_client('cloudtrail', region_name=region)
            response = cloudtrail_client.list_trails()
            trail_arns = [trail['TrailARN'] for trail in response['Trails']]
            return {"status": "success", "trails": trail_arns}
        except Exception as e:
            print(f"Error listing CloudTrail trails: {e}")
            return {"status": "error", "message": str(e)}

    def _describe_trails(self, region: str):
        """Describes all CloudTrail trails in a specified region."""
        print(f"CloudTrailAgent: Describing trails in region {region}...")
        try:
            cloudtrail_client = AWSConnector.get_client('cloudtrail', region_name=region)
            response = cloudtrail_client.describe_trails()
            trails_info = []
            for trail in response['trailList']:
                trails_info.append({
                    'Name': trail.get('Name'),
                    'TrailARN': trail.get('TrailARN'),
                    'HomeRegion': trail.get('HomeRegion'),
                    'IsMultiRegionTrail': trail.get('IsMultiRegionTrail'),
                    'S3BucketName': trail.get('S3BucketName'),
                })
            return {"status": "success", "trails_details": trails_info}
        except Exception as e:
            print(f"Error describing CloudTrail trails: {e}")
            return {"status": "error", "message": str(e)}