from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class S3Agent(BaseAgent):
    """
    An agent specialized in handling AWS S3 tasks.

    This agent provides methods for interacting with the S3 service, such as
    listing, creating, and deleting buckets. It utilizes the AWSConnector to get the necessary S3 client.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to S3.

        This method routes the command to the appropriate private method.
        For destructive actions like 'delete_bucket', it includes a confirmation step.

        Args:
            command (str): The specific action (e.g., 'list_buckets').
            **kwargs: A dictionary of arguments for the command.

        Returns:
            dict: The result of the command execution.
        
        Raises:
            NotImplementedError: If the command is not supported.
        """
        if command == 'list_buckets':
            return self._list_buckets(**kwargs)
        elif command == 'create_bucket':
            return self._create_bucket(**kwargs)
        elif command == 'smart_create_bucket':
            bucket_name = kwargs.get('bucket_name', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create an S3 bucket '{bucket_name}' in {region} with best practices (block public access, default encryption).")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_bucket(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create S3 bucket command cancelled by user."}
        elif command == 'delete_bucket':
            bucket_name = kwargs.get('bucket_name', '')
            print(f"WARNING: You are about to delete the S3 bucket '{bucket_name}'. This action is irreversible.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._delete_bucket(**kwargs)
            else:
                return {"status": "cancelled", "message": "Delete bucket command cancelled by user."}
        elif command == 'analyze_bucket_cost_optimization':
            bucket_name = kwargs.get('bucket_name', '')
            region = kwargs.get('region', '') # Assuming region is also needed for analysis
            print(f"You are about to analyze S3 bucket '{bucket_name}' for cost optimization.")
            confirm = input("This will retrieve bucket size and suggest lifecycle policies. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._analyze_bucket_cost_optimization(**kwargs)
            else:
                return {"status": "cancelled", "message": "Analyze bucket cost optimization command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by S3Agent.")
    def _list_buckets(self):
        """
        Lists all S3 buckets.
        """
        print("S3Agent: Listing all buckets...")
        try:
            s3_client = AWSConnector.get_client('s3')
            response = s3_client.list_buckets()
            
            bucket_names = [bucket['Name'] for bucket in response['Buckets']]
            
            if not bucket_names:
                return {"status": "success", "message": "No S3 buckets found."}
            
            return {"status": "success", "buckets": bucket_names}
        except Exception as e:
            print(f"Error listing S3 buckets: {e}")
            return {"status": "error", "message": str(e)}

    def _create_bucket(self, bucket_name: str, region: str):
        """
        Creates a new S3 bucket.
        """
        print(f"S3Agent: Creating bucket '{bucket_name}' in region {region}...")
        try:
            s3_client = AWSConnector.get_client('s3', region_name=region)
            # The LocationConstraint must be specified for regions other than us-east-1
            if region == "us-east-1":
                s3_client.create_bucket(Bucket=bucket_name)
            else:
                            location = {'LocationConstraint': region}
                            s3_client.create_bucket(Bucket=bucket_name, CreateBucketConfiguration=location)
            return {"status": "success", "message": f"Bucket '{bucket_name}' created successfully in region {region}."}
        except Exception as e:
            print(f"Error creating S3 bucket: {e}")
            return {"status": "error", "message": str(e)}

    def _delete_bucket(self, bucket_name: str):
        """
        Deletes an S3 bucket.
        """
        print(f"S3Agent: Deleting bucket '{bucket_name}'...")
        try:
            # Note: For this to succeed, the bucket must be empty.
            s3_client = AWSConnector.get_client('s3')
            s3_client.delete_bucket(Bucket=bucket_name)
            return {"status": "success", "message": f"Bucket '{bucket_name}' deleted successfully."}
        except Exception as e:
            print(f"Error deleting S3 bucket: {e}")
            return {"status": "error", "message": str(e)}

    def _smart_create_bucket(self, bucket_name: str, region: str, enable_logging: bool = False, logging_bucket_name: str = None):
        """
        Creates an S3 bucket with best practices:
        - Enforces block public access settings.
        - Enables default encryption.
        - Optionally enables server access logging.
        """
        print(f"S3Agent: Smart creating bucket '{bucket_name}' in region {region}...")
        try:
                        s3_client = AWSConnector.get_client('s3', region_name=region)
            
                        # 1. Create the bucket
                        if region == "us-east-1":
                            s3_client.create_bucket(Bucket=bucket_name)
                        else:
                            location = {'LocationConstraint': region}
                            s3_client.create_bucket(Bucket=bucket_name, CreateBucketConfiguration=location)
                        print(f"Bucket '{bucket_name}' created.")
            
                        # 2. Enforce Block Public Access
                        s3_client.put_public_access_block(
                            Bucket=bucket_name,
                            PublicAccessBlockConfiguration={
                                'BlockPublicAcls': True,
                                'IgnorePublicAcls': True,
                                'BlockPublicPolicy': True,
                                'RestrictPublicBuckets': True
                            }
                        )
                        print("Block Public Access enforced.")
            
                        # 3. Enable Default Encryption (SSE-S3)
                        s3_client.put_bucket_encryption(
                            Bucket=bucket_name,
                            ServerSideEncryptionConfiguration={
                                'Rules': [
                                    {
                                        'ApplyServerSideEncryptionByDefault': {
                                            'SSEAlgorithm': 'AES256'
                                        }
                                    },
                                ]
                            }
                        )
                        print("Default encryption (SSE-S3) enabled.")
            
                        # 4. Optionally enable server access logging
                        if enable_logging and logging_bucket_name:
                            # Check if logging bucket exists, if not, create it (simple version)
                            try:
                                s3_client.head_bucket(Bucket=logging_bucket_name)
                            except s3_client.exceptions.ClientError as e:
                                if e.response['Error']['Code'] == '404':
                                    print(f"Logging bucket '{logging_bucket_name}' not found. Creating it...")
                                    if region == "us-east-1":
                                        s3_client.create_bucket(Bucket=logging_bucket_name)
                                    else:
                                        location = {'LocationConstraint': region}
                                        s3_client.create_bucket(Bucket=logging_bucket_name, CreateBucketConfiguration=location)
                                    # Block public access and encryption for logging bucket too
                                    s3_client.put_public_access_block(
                                        Bucket=logging_bucket_name,
                                        PublicAccessBlockConfiguration={
                                            'BlockPublicAcls': True,
                                            'IgnorePublicAcls': True,
                                            'BlockPublicPolicy': True,
                                            'RestrictPublicBuckets': True
                                        }
                                    )
                                    s3_client.put_bucket_encryption(
                                        Bucket=logging_bucket_name,
                                        ServerSideEncryptionConfiguration={
                                            'Rules': [
                                                {
                                                    'ApplyServerSideEncryptionByDefault': {
                                                        'SSEAlgorithm': 'AES256'
                                                    }
                                                },
                                            ]
                                        }
                                    )
                                else:
                                    raise
            
                            s3_client.put_bucket_logging(
                                Bucket=bucket_name,
                                BucketLoggingStatus={
                                    'LoggingEnabled': {
                                        'TargetBucket': logging_bucket_name,
                                        'TargetPrefix': f'{bucket_name}/'
                                    }
                                }
                            )
                            print(f"Server access logging enabled, sending logs to '{logging_bucket_name}'.")
            
                        return {"status": "success", "message": f"Smart created bucket '{bucket_name}' with best practices in region {region}."}
        except Exception as e:
            print(f"Error during smart S3 bucket creation: {e}")
            return {"status": "error", "message": str(e)}

    def _analyze_bucket_cost_optimization(self, bucket_name: str):
        """
        Analyzes an S3 bucket for cost optimization opportunities, suggesting lifecycle policies.
        """
        print(f"S3Agent: Analyzing bucket '{bucket_name}' for cost optimization...")
        try:
            s3_client = AWSConnector.get_client('s3')
            
            # Get bucket size (simplified - head_bucket doesn't give size directly, requires list_objects_v2 or CloudWatch)
            # For a real implementation, you'd use CloudWatch metrics or S3 Inventory.
            # Here, we'll simulate based on object count for simplicity.
            object_count = 0
            total_size_bytes = 0
            try:
                response = s3_client.list_objects_v2(Bucket=bucket_name)
                if 'Contents' in response:
                    object_count = len(response['Contents'])
                    total_size_bytes = sum(obj['Size'] for obj in response['Contents'])
            except Exception as e:
                print(f"Warning: Could not get object details for bucket '{bucket_name}': {e}")
            
            suggestions = []
            if object_count > 1000 or total_size_bytes > (10 * 1024 * 1024 * 1024): # 10GB
                suggestions.append("Consider implementing a lifecycle policy to transition older objects to S3 Standard-IA or Glacier.")
                suggestions.append("For objects accessed infrequently, consider S3 Intelligent-Tiering.")
            
            # Check for existing lifecycle policies
            try:
                lifecycle_response = s3_client.get_bucket_lifecycle_configuration(Bucket=bucket_name)
                suggestions.append(f"Existing lifecycle policies: {lifecycle_response['Rules']}")
            except s3_client.exceptions.ClientError as e:
                if e.response['Error']['Code'] == 'NoSuchLifecycleConfiguration':
                    suggestions.append("No lifecycle policy found. Implementing one can save costs.")
                else:
                    raise
            
            if not suggestions:
                suggestions.append("No immediate cost optimization suggestions based on current analysis.")
            
            return {"status": "success", "message": f"Cost optimization analysis for bucket '{bucket_name}'.", "suggestions": suggestions, "object_count": object_count, "total_size_bytes": total_size_bytes}
        except Exception as e:
            print(f"Error analyzing S3 bucket for cost optimization: {e}")
            return {"status": "error", "message": str(e)}
            
