from .base_agent import BaseAgent
from ..aws_connector import AWSConnector
import time

class RedshiftAgent(BaseAgent):
    """
    An agent specialized in handling AWS Redshift tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Redshift.
        """
        if command == 'smart_create_cluster':
            cluster_identifier = kwargs.get('cluster_identifier', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create a Redshift cluster '{cluster_identifier}' in {region}.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_cluster(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create Redshift cluster command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by RedshiftAgent.")

    def _smart_create_cluster(self, region: str, cluster_identifier: str, node_type: str, master_username: str, master_user_password: str, node_count: int = 1, vpc_security_group_ids: list = None, vpc_id: str = None):
        """
        Creates a Redshift cluster.
        Automatically creates a security group if not provided.
        """
        print(f"RedshiftAgent: Smart creating cluster '{cluster_identifier}' in region {region}...")
        try:
            redshift_client = AWSConnector.get_client('redshift', region_name=region)
            ec2_client = AWSConnector.get_client('ec2', region_name=region)

            # 1. Determine VPC (if needed for SG creation)
            if vpc_id is None:
                vpcs = ec2_client.describe_vpcs(Filters=[{'Name': 'isDefault', 'Values': ['true']}]).get('Vpcs', [])
                if vpcs:
                    vpc_id = vpcs[0]['VpcId']
                else:
                    return {"status": "error", "message": "No default VPC found and no VPC ID provided for security group creation."}

            # 2. Create Security Group for Redshift if not provided
            if vpc_security_group_ids is None:
                sg_name = f'{cluster_identifier}-redshift-sg'
                sg_description = f'Security group for Redshift cluster {cluster_identifier}'
                try:
                    sg_response = ec2_client.create_security_group(GroupName=sg_name, Description=sg_description, VpcId=vpc_id)
                    sg_id = sg_response['GroupId']
                    # Allow Redshift default port (5439) from anywhere
                    ec2_client.authorize_security_group_ingress(
                        GroupId=sg_id,
                        IpPermissions=[
                            {'IpProtocol': 'tcp', 'FromPort': 5439, 'ToPort': 5439, 'IpRanges': [{'CidrIp': '0.0.0.0/0'}]}
                        ]
                    )
                    vpc_security_group_ids = [sg_id]
                    print(f"Created Security Group: {sg_id}")
                except ec2_client.exceptions.ClientError as e:
                    if 'already exists' in str(e):
                        print(f"Security Group '{sg_name}' already exists. Attempting to use existing.")
                        sg_details = ec2_client.describe_security_groups(Filters=[{'Name': 'group-name', 'Values': [sg_name]}, {'Name': 'vpc-id', 'Values': [vpc_id]}]).get('SecurityGroups', [])
                        if sg_details:
                            vpc_security_group_ids = [sg_details[0]['GroupId']]
                        else:
                            raise Exception(f"Could not find existing security group '{sg_name}'.")
                    else:
                        raise

            cluster_params = {
                'ClusterIdentifier': cluster_identifier,
                'NodeType': node_type,
                'MasterUsername': master_username,
                'MasterUserPassword': master_user_password,
                'NumberOfNodes': node_count,
                'PubliclyAccessible': True, # For simplicity
                'Tags': [
                    {'Key': 'Name', 'Value': cluster_identifier}
                ],
                'VpcSecurityGroupIds': vpc_security_group_ids
            }

            response = redshift_client.create_cluster(**cluster_params)
            cluster_arn = response['Cluster']['ClusterIdentifier'] # ARN is not directly returned, using identifier

            # Wait for cluster to be available
            print("Waiting for Redshift cluster to become available... (this may take several minutes)")
            waiter = redshift_client.get_waiter('cluster_available')
            waiter.wait(ClusterIdentifier=cluster_identifier)

            return {"status": "success", "message": f"Redshift cluster '{cluster_identifier}' created successfully.", "cluster_identifier": cluster_identifier, "security_group_ids": vpc_security_group_ids}
        except Exception as e:
            print(f"Error during smart Redshift cluster creation: {e}")
            return {"status": "error", "message": str(e)}