from ..base_gcp_agent import BaseGCPAgent
from ..gcp_connector import GCPConnector
from google.cloud import compute_v1

class ComputeAgent(BaseGCPAgent):
    """
    An agent specialized in handling GCP Compute Engine tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Compute Engine.
        """
        if command == 'list_instances':
            return self._list_instances(**kwargs)
        elif command == 'start_instance':
            return self._start_instance(**kwargs)
        elif command == 'stop_instance':
            return self._stop_instance(**kwargs)
        elif command == 'delete_instance':
            instance_name = kwargs.get('instance_name', '')
            print(f"WARNING: You are about to delete Compute Engine instance '{instance_name}'. This action is irreversible.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._delete_instance(**kwargs)
            else:
                return {"status": "cancelled", "message": "Delete instance command cancelled by user."}
        elif command == 'smart_create_instance':
            instance_name = kwargs.get('instance_name', '')
            print(f"You are about to smart-create a Compute Engine instance '{instance_name}'.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_instance(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create instance command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by ComputeAgent.")

    def _list_instances(self, project_id: str, zone: str):
        """
        Lists all Compute Engine instances in a specified zone.
        """
        print(f"ComputeAgent: Listing instances in project '{project_id}', zone '{zone}'...")
        try:
            client = GCPConnector.get_compute_client()
            request = compute_v1.ListInstancesRequest(project=project_id, zone=zone)
            response = client.list(request=request)
            
            instances = []
            for instance in response.items:
                instances.append({
                    'name': instance.name,
                    'status': instance.status,
                    'machine_type': instance.machine_type.split('/')[-1],
                    'internal_ip': instance.network_interfaces[0].network_ip if instance.network_interfaces else 'N/A',
                    'external_ip': instance.network_interfaces[0].access_configs[0].nat_ip if instance.network_interfaces and instance.network_interfaces[0].access_configs else 'N/A',
                })
            return {"status": "success", "instances": instances}
        except Exception as e:
            print(f"Error listing instances: {e}")
            return {"status": "error", "message": str(e)}

    def _start_instance(self, project_id: str, zone: str, instance_name: str):
        """
        Starts a specified Compute Engine instance.
        """
        print(f"ComputeAgent: Starting instance '{instance_name}' in project '{project_id}', zone '{zone}'...")
        try:
            client = GCPConnector.get_compute_client()
            request = compute_v1.StartInstanceRequest(project=project_id, zone=zone, instance=instance_name)
            operation = client.start(request=request)
            # Wait for the operation to complete
            operation.wait()
            return {"status": "success", "message": f"Instance '{instance_name}' started successfully."}
        except Exception as e:
            print(f"Error starting instance: {e}")
            return {"status": "error", "message": str(e)}

    def _stop_instance(self, project_id: str, zone: str, instance_name: str):
        """
        Stops a specified Compute Engine instance.
        """
        print(f"ComputeAgent: Stopping instance '{instance_name}' in project '{project_id}', zone '{zone}'...")
        try:
            client = GCPConnector.get_compute_client()
            request = compute_v1.StopInstanceRequest(project=project_id, zone=zone, instance=instance_name)
            operation = client.stop(request=request)
            # Wait for the operation to complete
            operation.wait()
            return {"status": "success", "message": f"Instance '{instance_name}' stopped successfully."}
        except Exception as e:
            print(f"Error stopping instance: {e}")
            return {"status": "error", "message": str(e)}

    def _delete_instance(self, project_id: str, zone: str, instance_name: str):
        """
        Deletes a specified Compute Engine instance.
        """
        print(f"ComputeAgent: Deleting instance '{instance_name}' in project '{project_id}', zone '{zone}'...")
        try:
            client = GCPConnector.get_compute_client()
            request = compute_v1.DeleteInstanceRequest(project=project_id, zone=zone, instance=instance_name)
            operation = client.delete(request=request)
            # Wait for the operation to complete
            operation.wait()
            return {"status": "success", "message": f"Instance '{instance_name}' deleted successfully."}
        except Exception as e:
            print(f"Error deleting instance: {e}")
            return {"status": "error", "message": str(e)}

    def _smart_create_instance(self, project_id: str, zone: str, instance_name: str, machine_type: str = 'e2-medium', image_project: str = 'debian-cloud', image_family: str = 'debian-11'):
        """
        Creates a Compute Engine instance with sensible defaults.
        """
        print(f"ComputeAgent: Smart creating instance '{instance_name}' in project '{project_id}', zone '{zone}'...")
        try:
            client = GCPConnector.get_compute_client()

            # Get the latest image from the specified family
            image_client = compute_v1.ImagesClient(credentials=GCPConnector.get_credentials())
            image_request = compute_v1.GetFromFamilyImageRequest(project=image_project, family=image_family)
            image = image_client.get_from_family(request=image_request)

            # Configure the machine type
            machine_type_url = f"zones/{zone}/machineTypes/{machine_type}"

            # Create a disk from the image
            disk = compute_v1.AttachedDisk()
            disk.initialize_params = compute_v1.AttachedDiskInitializeParams()
            disk.initialize_params.source_image = image.self_link
            disk.auto_delete = True
            disk.boot = True

            # Configure network interface
            network_interface = compute_v1.NetworkInterface()
            network_interface.name = 'nic0'
            network_interface.network = f"projects/{project_id}/global/networks/default"
            access_config = compute_v1.AccessConfig()
            access_config.name = 'External NAT'
            access_config.type_ = 'ONE_TO_ONE_NAT'
            network_interface.access_configs = [access_config]

            # Create the instance
            instance = compute_v1.Instance()
            instance.name = instance_name
            instance.machine_type = machine_type_url
            instance.disks = [disk]
            instance.network_interfaces = [network_interface]

            request = compute_v1.InsertInstanceRequest(
                project=project_id,
                zone=zone,
                instance_resource=instance,
            )

            operation = client.insert(request=request)
            operation.wait()

            return {"status": "success", "message": f"Compute Engine instance '{instance_name}' created successfully."}
        except Exception as e:
            print(f"Error creating instance: {e}")
            return {"status": "error", "message": str(e)}