from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class ELBAgent(BaseAgent):
    """
    An agent specialized in handling AWS Elastic Load Balancing (ELB) tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to ELB.
        """
        if command == 'list_load_balancers':
            return self._list_load_balancers(**kwargs)
        elif command == 'create_load_balancer':
            return self._create_load_balancer(**kwargs)
        elif command == 'create_target_group':
            return self._create_target_group(**kwargs)
        elif command == 'register_targets':
            return self._register_targets(**kwargs)
        elif command == 'create_listener':
            return self._create_listener(**kwargs)
        elif command == 'smart_create_alb':
            name = kwargs.get('name', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create an Application Load Balancer '{name}' in {region}.")
            confirm = input("This will create an ALB, Target Group, and Listener. It may also create a VPC and Security Group. Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_alb(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create ALB command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by ELBAgent.")

    def _smart_create_alb(self, region: str, name: str, vpc_id: str = None, subnet_ids: list = None, security_group_ids: list = None):
        """
        Creates an Application Load Balancer, a target group, and a listener.
        Automatically handles VPC and subnet creation if not provided.
        """
        print(f"ELBAgent: Smart creating ALB '{name}' in region {region}...")
        try:
            ec2_client = AWSConnector.get_client('ec2', region_name=region)
            
            # 1. Determine VPC and Subnets
            if vpc_id is None or subnet_ids is None:
                print("VPC ID or Subnet IDs not provided. Attempting to find/create VPC...")
                # Instantiate AgentManager to use VPC workflow
                agent_manager = AgentManager() 
                vpc_result = agent_manager._workflow_smart_create_vpc(region=region, name=f'{name}-alb-vpc')
                if vpc_result['status'] != 'success':
                    return vpc_result
                vpc_id = vpc_result['outputs']['vpc_id']
                subnet_ids = [vpc_result['outputs']['public_subnet_id'], vpc_result['outputs']['private_subnet_id']]
                print(f"Using VPC: {vpc_id} with subnets: {subnet_ids}")

            # 2. Create Security Group for ALB if not provided
            if security_group_ids is None:
                sg_name = f'{name}-alb-sg'
                sg_description = f'Security group for ALB {name}'
                try:
                    sg_response = ec2_client.create_security_group(GroupName=sg_name, Description=sg_description, VpcId=vpc_id)
                    sg_id = sg_response['GroupId']
                    # Allow HTTP traffic from anywhere
                    ec2_client.authorize_security_group_ingress(
                        GroupId=sg_id,
                        IpPermissions=[
                            {'IpProtocol': 'tcp', 'FromPort': 80, 'ToPort': 80, 'IpRanges': [{'CidrIp': '0.0.0.0/0'}]}
                        ]
                    )
                    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:
                            security_group_ids = [sg_details[0]['GroupId']]
                        else:
                            raise Exception(f"Could not find existing security group '{sg_name}'.")
                    else:
                        raise

            # 3. Create ALB
            alb_response = self._create_load_balancer(region=region, name=name, subnets=subnet_ids, security_groups=security_group_ids)
            if alb_response['status'] != 'success':
                return alb_response
            alb_arn = alb_response['load_balancer']['LoadBalancerArn']
            alb_dns = alb_response['load_balancer']['DNSName']
            print(f"ALB '{name}' created with DNS: {alb_dns}")

            # 4. Create Target Group
            tg_name = f'{name}-tg'
            tg_response = self._create_target_group(region=region, name=tg_name, vpc_id=vpc_id)
            if tg_response['status'] != 'success':
                return tg_response
            tg_arn = tg_response['target_group']['TargetGroupArn']
            print(f"Target Group '{tg_name}' created.")

            # 5. Create Listener
            listener_response = self._create_listener(region=region, load_balancer_arn=alb_arn, target_group_arn=tg_arn)
            if listener_response['status'] != 'success':
                return listener_response
            print("Listener created.")

            return {"status": "success", "message": f"Smart created ALB '{name}' with DNS '{alb_dns}', Target Group '{tg_name}', and Listener.", "alb_arn": alb_arn, "alb_dns": alb_dns, "target_group_arn": tg_arn, "vpc_id": vpc_id, "subnet_ids": subnet_ids, "security_group_ids": security_group_ids}
        except Exception as e:
            print(f"Error during smart ALB creation: {e}")
            return {"status": "error", "message": str(e)}