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

class VPCAgent(BaseAgent):
    """
    An agent specialized in handling AWS VPC tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to VPC.
        """
        if command == 'create_vpc':
            pass # Placeholder for code
        elif command == 'create_subnet':
            pass # Placeholder for code
        elif command == 'smart_create_vpc':
            print(f"You are about to create a complete VPC environment named '{kwargs.get('name')}'.")
            confirm = input("This will create a VPC, Subnets, IGW, NAT GW, and Route Tables. Are you sure? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_vpc(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart VPC creation cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by VPCAgent.")

    def _create_vpc(self, region: str, cidr_block: str):
        # ... (code unchanged)
        pass

    def _create_subnet(self, region: str, vpc_id: str, cidr_block: str):
        # ... (code unchanged)
        pass

    def _smart_create_vpc(self, region: str, name: str, cidr_block: str = '10.0.0.0/16', num_public_subnets: int = 1, num_private_subnets: int = 1):
        """
        Creates a complete, best-practice VPC environment with a specified number of public and private subnets.
        """
        print(f"VPCAgent: Beginning Smart Create for VPC environment '{name}' in region {region}...")
        try:
            ec2_client = AWSConnector.get_client('ec2', region_name=region)
            resources = {}

            # 1. Create VPC
            vpc_res = ec2_client.create_vpc(CidrBlock=cidr_block)
            vpc_id = vpc_res['Vpc']['VpcId']
            resources['vpc_id'] = vpc_id
            ec2_client.create_tags(Resources=[vpc_id], Tags=[{'Key': 'Name', 'Value': name}])
            print(f"Created VPC: {vpc_id}")

            # Enable DNS Hostnames and Resolution
            ec2_client.modify_vpc_attribute(VpcId=vpc_id, EnableDnsSupport={'Value': True})
            ec2_client.modify_vpc_attribute(VpcId=vpc_id, EnableDnsHostnames={'Value': True})

            # Get available Availability Zones
            azs_response = ec2_client.describe_availability_zones(Filters=[{'Name': 'state', 'Values': ['available']}])
            available_azs = [az['ZoneName'] for az in azs_response['AvailabilityZones']]
            if not available_azs:
                return {"status": "error", "message": "No available Availability Zones in the specified region."}

            # 2. Create Internet Gateway
            igw_res = ec2_client.create_internet_gateway()
            igw_id = igw_res['InternetGateway']['InternetGatewayId']
            resources['internet_gateway_id'] = igw_id
            ec2_client.attach_internet_gateway(InternetGatewayId=igw_id, VpcId=vpc_id)
            print(f"Created and attached Internet Gateway: {igw_id}")

            # 3. Create Public Subnets
            public_subnet_ids = []
            public_route_table_id = None
            for i in range(num_public_subnets):
                if i >= len(available_azs):
                    print(f"Warning: Not enough AZs for {num_public_subnets} public subnets. Creating {len(available_azs)} public subnets.")
                    break
                subnet_cidr = f'10.0.{i}.0/24'
                subnet_res = ec2_client.create_subnet(VpcId=vpc_id, CidrBlock=subnet_cidr, AvailabilityZone=available_azs[i])
                subnet_id = subnet_res['Subnet']['SubnetId']
                public_subnet_ids.append(subnet_id)
                ec2_client.create_tags(Resources=[subnet_id], Tags=[{'Key': 'Name', 'Value': f'{name}-public-{i+1}'}])
                ec2_client.modify_subnet_attribute(SubnetId=subnet_id, MapPublicIpOnLaunch={'Value': True})
                print(f"Created Public Subnet {i+1}: {subnet_id} in {available_azs[i]}")

                # Create/Associate Public Route Table for the first public subnet
                if public_route_table_id is None:
                    public_rt_res = ec2_client.create_route_table(VpcId=vpc_id)
                    public_route_table_id = public_rt_res['RouteTable']['RouteTableId']
                    ec2_client.create_route(RouteTableId=public_route_table_id, DestinationCidrBlock='0.0.0.0/0', GatewayId=igw_id)
                    resources['public_route_table_id'] = public_route_table_id
                    print(f"Created Public Route Table: {public_route_table_id}")
                ec2_client.associate_route_table(RouteTableId=public_route_table_id, SubnetId=subnet_id)
                print(f"Associated Public Route Table with public subnet {i+1}")
            resources['public_subnet_ids'] = public_subnet_ids

            # 4. Create NAT Gateway (in the first public subnet)
            nat_gw_id = None
            if public_subnet_ids:
                eip_res = ec2_client.allocate_address(Domain='vpc')
                eip_alloc_id = eip_res['AllocationId']
                resources['nat_gateway_eip_allocation_id'] = eip_alloc_id
                print(f"Allocated Elastic IP for NAT Gateway: {eip_res['PublicIp']}")

                nat_res = ec2_client.create_nat_gateway(SubnetId=public_subnet_ids[0], AllocationId=eip_alloc_id)
                nat_gw_id = nat_res['NatGateway']['NatGatewayId']
                resources['nat_gateway_id'] = nat_gw_id
                print("Waiting for NAT Gateway to become available...")
                waiter = ec2_client.get_waiter('nat_gateway_available')
                waiter.wait(NatGatewayIds=[nat_gw_id])
                print(f"NAT Gateway available: {nat_gw_id}")
            else:
                print("No public subnets created, skipping NAT Gateway creation.")

            # 5. Create Private Subnets and Route Tables
            private_subnet_ids = []
            for i in range(num_private_subnets):
                if i >= len(available_azs):
                    print(f"Warning: Not enough AZs for {num_private_subnets} private subnets. Creating {len(available_azs)} private subnets.")
                    break
                subnet_cidr = f'10.0.{i + num_public_subnets}.0/24'
                subnet_res = ec2_client.create_subnet(VpcId=vpc_id, CidrBlock=subnet_cidr, AvailabilityZone=available_azs[i])
                subnet_id = subnet_res['Subnet']['SubnetId']
                private_subnet_ids.append(subnet_id)
                ec2_client.create_tags(Resources=[subnet_id], Tags=[{'Key': 'Name', 'Value': f'{name}-private-{i+1}'}])
                print(f"Created Private Subnet {i+1}: {subnet_id} in {available_azs[i]}")

                # Create Private Route Table and associate with private subnet
                if nat_gw_id:
                    private_rt_res = ec2_client.create_route_table(VpcId=vpc_id)
                    private_rt_id = private_rt_res['RouteTable']['RouteTableId']
                    ec2_client.create_route(RouteTableId=private_rt_id, DestinationCidrBlock='0.0.0.0/0', NatGatewayId=nat_gw_id)
                    ec2_client.associate_route_table(RouteTableId=private_rt_id, SubnetId=subnet_id)
                    resources[f'private_route_table_id_{i+1}'] = private_rt_id
                    print(f"Created Private Route Table and associated with private subnet {i+1}: {private_rt_id}")
                else:
                    print(f"No NAT Gateway, skipping private route table for private subnet {i+1}.")
            resources['private_subnet_ids'] = private_subnet_ids

            return {"status": "success", "message": f"Successfully created VPC environment '{name}'.", "created_resources": resources}

        except Exception as e:
            print(f"Error during smart VPC creation: {e}")
            # Basic cleanup logic - could be more robust
            print("Attempting to clean up created resources...")
            if 'nat_gateway_id' in resources:
                ec2_client.delete_nat_gateway(NatGatewayId=resources['nat_gateway_id'])
            if 'internet_gateway_id' in resources:
                ec2_client.detach_internet_gateway(InternetGatewayId=resources['internet_gateway_id'], VpcId=resources['vpc_id'])
                ec2_client.delete_internet_gateway(InternetGatewayId=resources['internet_gateway_id'])
            if 'nat_gateway_eip_allocation_id' in resources:
                ec2_client.release_address(AllocationId=resources['nat_gateway_eip_allocation_id'])
            if 'public_subnet_ids' in resources:
                for subnet_id in resources['public_subnet_ids']:
                    ec2_client.delete_subnet(SubnetId=subnet_id)
            if 'private_subnet_ids' in resources:
                for subnet_id in resources['private_subnet_ids']:
                    ec2_client.delete_subnet(SubnetId=subnet_id)
            if 'vpc_id' in resources:
                ec2_client.delete_vpc(VpcId=resources['vpc_id'])
            return {"status": "error", "message": str(e)}