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

class EFSAgent(BaseAgent):
    """
    An agent specialized in handling AWS Elastic File System (EFS) tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to EFS.
        """
        if command == 'smart_create_file_system':
            name = kwargs.get('name', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create an EFS file system '{name}' in {region}.")
            confirm = input("This will create an EFS file system and a mount target. 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_file_system(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create EFS file system command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by EFSAgent.")

    def _smart_create_file_system(self, region: str, name: str, subnet_id: str = None, vpc_id: str = None, performance_mode: str = 'generalPurpose', throughput_mode: str = 'bursting', security_group_id: str = None):
        """
        Creates an EFS file system and a mount target in a specified subnet.
        Automatically handles VPC and subnet creation if not provided.
        """
        print(f"EFSAgent: Smart creating file system '{name}' in region {region}...")
        try:
            efs_client = AWSConnector.get_client('efs', region_name=region)
            ec2_client = AWSConnector.get_client('ec2', region_name=region)

            # 1. Determine VPC and Subnet
            if vpc_id is None or subnet_id is None:
                print("VPC ID or Subnet ID 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}-efs-vpc')
                if vpc_result['status'] != 'success':
                    return vpc_result
                vpc_id = vpc_result['outputs']['vpc_id']
                # Use a public subnet for simplicity, or a private one if specified
                subnet_id = vpc_result['outputs']['public_subnet_id'] 
                print(f"Using VPC: {vpc_id} with subnet: {subnet_id}")

            # 2. Create EFS File System
            fs_response = efs_client.create_file_system(
                CreationToken=f'{name}-{int(time.time())}',
                PerformanceMode=performance_mode,
                ThroughputMode=throughput_mode,
                Tags=[
                    {'Key': 'Name', 'Value': name},
                ]
            )
            file_system_id = fs_response['FileSystemId']
            print(f"EFS File System '{file_system_id}' created.")

            # Wait for file system to be available
            waiter = efs_client.get_waiter('file_system_available')
            waiter.wait(FileSystemId=file_system_id)
            print(f"EFS File System '{file_system_id}' is now available.")

            # 3. Create Security Group for Mount Target if not provided
            if security_group_id is None:
                sg_name = f'{name}-efs-sg'
                sg_description = f'Security group for EFS mount target {name}'
                try:
                    sg_response = ec2_client.create_security_group(GroupName=sg_name, Description=sg_description, VpcId=vpc_id)
                    security_group_id = sg_response['GroupId']
                    # Authorize NFS ingress from within the VPC
                    ec2_client.authorize_security_group_ingress(
                        GroupId=security_group_id,
                        IpPermissions=[
                            {'IpProtocol': 'tcp', 'FromPort': 2049, 'ToPort': 2049, 'IpRanges': [{'CidrIp': '0.0.0.0/0'}]} # Allow NFS from anywhere for simplicity, refine in production
                        ]
                    )
                    print(f"Created Security Group: {security_group_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_id = sg_details[0]['GroupId']
                        else:
                            raise Exception(f"Could not find existing security group '{sg_name}'.")
                    else:
                        raise

            # 4. Create Mount Target
            mt_response = efs_client.create_mount_target(
                FileSystemId=file_system_id,
                SubnetId=subnet_id,
                SecurityGroups=[security_group_id]
            )
            mount_target_id = mt_response['MountTargetId']
            print(f"EFS Mount Target '{mount_target_id}' created in subnet '{subnet_id}'.")

            # Wait for mount target to be available
            waiter = efs_client.get_waiter('mount_target_available')
            waiter.wait(MountTargetId=mount_target_id)
            print(f"EFS Mount Target '{mount_target_id}' is now available.")

            return {"status": "success", "message": f"EFS file system '{name}' ({file_system_id}) created with mount target '{mount_target_id}'.", "file_system_id": file_system_id, "mount_target_id": mount_target_id, "security_group_id": security_group_id, "vpc_id": vpc_id, "subnet_id": subnet_id}
        except Exception as e:
            print(f"Error during smart EFS file system creation: {e}")
            return {"status": "error", "message": str(e)}