from .base_agent import BaseAgent
from ..aws_connector import AWSConnector

class EBSAgent(BaseAgent):
    """
    An agent specialized in handling AWS Elastic Block Store (EBS) tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to EBS.
        """
        if command == 'smart_create_and_attach_volume':
            volume_name = kwargs.get('volume_name', '')
            instance_id = kwargs.get('instance_id', '')
            region = kwargs.get('region', '')
            print(f"You are about to smart-create and attach an EBS volume '{volume_name}' to instance '{instance_id}' in {region}.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_and_attach_volume(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create and Attach EBS Volume command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by EBSAgent.")

    def _smart_create_and_attach_volume(self, region: str, volume_name: str, instance_id: str, size: int = 10, volume_type: str = 'gp2'):
        """
        Creates a new EBS volume and attaches it to a specified EC2 instance.
        """
        print(f"EBSAgent: Smart creating and attaching volume '{volume_name}' to instance '{instance_id}' in region {region}...")
        try:
            ec2_client = AWSConnector.get_client('ec2', region_name=region)

            # 1. Get instance details to find Availability Zone
            instance_response = ec2_client.describe_instances(InstanceIds=[instance_id])
            reservations = instance_response['Reservations']
            if not reservations or not reservations[0]['Instances']:
                return {"status": "error", "message": f"EC2 instance '{instance_id}' not found in region {region}."}
            instance_az = reservations[0]['Instances'][0]['Placement']['AvailabilityZone']
            print(f"Instance '{instance_id}' is in Availability Zone: {instance_az}")

            # 2. Create EBS Volume
            volume_response = ec2_client.create_volume(
                AvailabilityZone=instance_az,
                Size=size,
                VolumeType=volume_type,
                TagSpecifications=[
                    {
                        'ResourceType': 'volume',
                        'Tags': [
                            {'Key': 'Name', 'Value': volume_name}
                        ]
                    },
                ]
            )
            volume_id = volume_response['VolumeId']
            print(f"EBS Volume '{volume_id}' created.")

            # Wait for volume to be available
            waiter = ec2_client.get_waiter('volume_available')
            waiter.wait(VolumeIds=[volume_id])
            print(f"EBS Volume '{volume_id}' is now available.")

            # 3. Attach Volume to Instance
            # Determine a suitable device name (e.g., /dev/sdf, /dev/sdg, etc.)
            # This is a simplified approach; in a real scenario, you might need to check existing attachments.
            device_name = '/dev/sdf' # Common default for Linux instances
            attach_response = ec2_client.attach_volume(
                Device=device_name,
                InstanceId=instance_id,
                VolumeId=volume_id
            )
            print(f"EBS Volume '{volume_id}' attached to '{instance_id}' as '{device_name}'.")

            return {"status": "success", "message": f"EBS volume '{volume_name}' ({volume_id}) created and attached to instance '{instance_id}'.", "volume_id": volume_id}
        except Exception as e:
            print(f"Error during smart EBS volume creation and attachment: {e}")
            return {"status": "error", "message": str(e)}