from ..base_azure_agent import BaseAzureAgent
from ..azure_connector import AzureConnector
from azure.mgmt.compute.models import VirtualMachine, HardwareProfile, OSProfile, StorageProfile, NetworkProfile, ImageReference, DataDisk, DiskCreateOptionTypes
from azure.mgmt.network.models import VirtualNetwork, Subnet, NetworkInterface, NetworkInterfaceIPConfiguration, PublicIPAddress, NetworkSecurityGroup, SecurityRule
from msrestazure.azure_exceptions import CloudError
import time

class VMAgent(BaseAzureAgent):
    """
    An agent specialized in handling Azure Virtual Machine tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Azure VMs.
        """
        if command == 'list_vms':
            return self._list_vms(**kwargs)
        elif command == 'start_vm':
            return self._start_vm(**kwargs)
        elif command == 'stop_vm':
            return self._stop_vm(**kwargs)
        elif command == 'delete_vm':
            vm_name = kwargs.get('vm_name', '')
            resource_group = kwargs.get('resource_group', '')
            print(f"WARNING: You are about to delete Azure VM '{vm_name}' in resource group '{resource_group}'. This action is irreversible.")
            confirm = input("Are you sure you want to proceed? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._delete_vm(**kwargs)
            else:
                return {"status": "cancelled", "message": "Delete VM command cancelled by user."}
        elif command == 'smart_create_vm':
            vm_name = kwargs.get('vm_name', '')
            resource_group = kwargs.get('resource_group', '')
            print(f"You are about to smart-create an Azure VM '{vm_name}' in resource group '{resource_group}'.")
            confirm = input("This will create a new VM, VNet, Subnet, Public IP, and Network Interface. Are you sure? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_vm(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create VM command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by VMAgent.")

    def _list_vms(self, subscription_id: str, resource_group: str):
        """
        Lists all Virtual Machines in a specified resource group.
        """
        print(f"VMAgent: Listing VMs in subscription '{subscription_id}', resource group '{resource_group}'...")
        try:
            compute_client = AzureConnector.get_compute_client()
            vms = compute_client.virtual_machines.list(resource_group)
            vm_details = []
            for vm in vms:
                vm_details.append({
                    'name': vm.name,
                    'id': vm.id,
                    'location': vm.location,
                    'hardware_profile': vm.hardware_profile.vm_size,
                    'provisioning_state': vm.provisioning_state,
                    'power_state': vm.instance_view.statuses[1].display_status if vm.instance_view and vm.instance_view.statuses else 'N/A'
                })
            return {"status": "success", "vms": vm_details}
        except CloudError as e:
            print(f"Azure Error listing VMs: {e.message}")
            return {"status": "error", "message": e.message}
        except Exception as e:
            print(f"Error listing VMs: {e}")
            return {"status": "error", "message": str(e)}

    def _start_vm(self, subscription_id: str, resource_group: str, vm_name: str):
        """
        Starts a specified Virtual Machine.
        """
        print(f"VMAgent: Starting VM '{vm_name}' in resource group '{resource_group}'...")
        try:
            compute_client = AzureConnector.get_compute_client()
            async_vm_start = compute_client.virtual_machines.begin_start(resource_group, vm_name)
            async_vm_start.wait()
            return {"status": "success", "message": f"VM '{vm_name}' started successfully."}
        except CloudError as e:
            print(f"Azure Error starting VM: {e.message}")
            return {"status": "error", "message": e.message}
        except Exception as e:
            print(f"Error starting VM: {e}")
            return {"status": "error", "message": str(e)}

    def _stop_vm(self, subscription_id: str, resource_group: str, vm_name: str):
        """
        Stops a specified Virtual Machine.
        """
        print(f"VMAgent: Stopping VM '{vm_name}' in resource group '{resource_group}'...")
        try:
            compute_client = AzureConnector.get_compute_client()
            async_vm_stop = compute_client.virtual_machines.begin_power_off(resource_group, vm_name)
            async_vm_stop.wait()
            return {"status": "success", "message": f"VM '{vm_name}' stopped successfully."}
        except CloudError as e:
            print(f"Azure Error stopping VM: {e.message}")
            return {"status": "error", "message": e.message}
        except Exception as e:
            print(f"Error stopping VM: {e}")
            return {"status": "error", "message": str(e)}

    def _delete_vm(self, subscription_id: str, resource_group: str, vm_name: str):
        """
        Deletes a specified Virtual Machine.
        """
        print(f"VMAgent: Deleting VM '{vm_name}' in resource group '{resource_group}'...")
        try:
            compute_client = AzureConnector.get_compute_client()
            async_vm_delete = compute_client.virtual_machines.begin_delete(resource_group, vm_name)
            async_vm_delete.wait()
            return {"status": "success", "message": f"VM '{vm_name}' deleted successfully."}
        except CloudError as e:
            print(f"Azure Error deleting VM: {e.message}")
            return {"status": "error", "message": e.message}
        except Exception as e:
            print(f"Error deleting VM: {e}")
            return {"status": "error", "message": str(e)}

    def _smart_create_vm(self, subscription_id: str, resource_group: str, location: str, vm_name: str, 
                          image_publisher: str = 'Canonical', image_offer: str = 'UbuntuServer', image_sku: str = '18.04-LTS', 
                          vm_size: str = 'Standard_B1s', admin_username: str = 'azureuser', admin_password: str = 'P@ssw0rd12345'):
        """
        Creates an Azure VM with sensible defaults, automatically handling network dependencies.
        """
        print(f"VMAgent: Smart creating VM '{vm_name}' in resource group '{resource_group}', location '{location}'...")
        try:
            compute_client = AzureConnector.get_compute_client()
            network_client = AzureConnector.get_network_client()

            # 1. Create Virtual Network (VNet) if it doesn't exist
            vnet_name = f'{vm_name}-vnet'
            subnet_name = f'{vm_name}-subnet'
            ip_config_name = f'{vm_name}-ipconfig'
            nic_name = f'{vm_name}-nic'
            public_ip_name = f'{vm_name}-public-ip'

            # Check if VNet exists
            try:
                vnet = network_client.virtual_networks.get(resource_group, vnet_name)
                print(f"Using existing VNet: {vnet_name}")
            except CloudError:
                print(f"Creating VNet: {vnet_name}")
                vnet_params = VirtualNetwork(
                    location=location,
                    address_space={'address_prefixes': ['10.0.0.0/16']},
                    subnets=[Subnet(name=subnet_name, address_prefix='10.0.0.0/24')]
                )
                async_vnet_creation = network_client.virtual_networks.begin_create_or_update(resource_group, vnet_name, vnet_params)
                vnet = async_vnet_creation.wait()
                print(f"VNet '{vnet_name}' created.")

            # 2. Create Public IP Address
            print(f"Creating Public IP: {public_ip_name}")
            public_ip_params = PublicIPAddress(
                location=location,
                public_ip_allocation_method=IpAllocationMethod.STATIC,
                sku={'name': 'Basic'}
            )
            async_public_ip_creation = network_client.public_ip_addresses.begin_create_or_update(resource_group, public_ip_name, public_ip_params)
            public_ip_address = async_public_ip_creation.wait()
            print(f"Public IP '{public_ip_name}' created.")

            # 3. Create Network Interface (NIC)
            print(f"Creating Network Interface: {nic_name}")
            nic_params = NetworkInterface(
                location=location,
                ip_configurations=[NetworkInterfaceIPConfiguration(
                    name=ip_config_name,
                    subnet={'id': vnet.subnets[0].id},
                    public_ip_address={'id': public_ip_address.id}
                )]
            )
            async_nic_creation = network_client.network_interfaces.begin_create_or_update(resource_group, nic_name, nic_params)
            network_interface = async_nic_creation.wait()
            print(f"Network Interface '{nic_name}' created.")

            # 4. Create Virtual Machine
            print(f"Creating VM: {vm_name}")
            vm_parameters = VirtualMachine(
                location=location,
                hardware_profile=HardwareProfile(vm_size=vm_size),
                os_profile=OSProfile(
                    computer_name=vm_name,
                    admin_username=admin_username,
                    admin_password=admin_password
                ),
                storage_profile=StorageProfile(
                    image_reference=ImageReference(
                        publisher=image_publisher,
                        offer=image_offer,
                        sku=image_sku,
                        version='latest'
                    ),
                    os_disk=DataDisk(
                        name=f'{vm_name}-osdisk',
                        create_option=DiskCreateOptionTypes.FROM_IMAGE,
                        managed_disk={'storage_account_type': 'Standard_LRS'}
                    )
                ),
                network_profile=NetworkProfile(
                    network_interfaces=[{'id': network_interface.id}]
                )
            )
            async_vm_creation = compute_client.virtual_machines.begin_create_or_update(resource_group, vm_name, vm_parameters)
            vm = async_vm_creation.wait()
            print(f"VM '{vm_name}' created.")

            return {"status": "success", "message": f"VM '{vm_name}' created successfully.", "vm_id": vm.id}
        except CloudError as e:
            print(f"Azure Error creating VM: {e.message}")
            return {"status": "error", "message": e.message}
        except Exception as e:
            print(f"Error creating VM: {e}")
            return {"status": "error", "message": str(e)}