#!/usr/bin/env python3
"""
Azure VM Launch Test
Launches a real Standard_B1ms instance on Azure free tier and runs a simple compute job.

WARNING: This test creates REAL cloud resources that cost money!
Set AZURE_REAL_TEST=true only when you explicitly want to test VM creation.
"""

import os
import sys
import time

# =============================================================================
# SAFETY FLAG - Set to True only when you want to launch real VMs
# =============================================================================
REAL_TEST = os.environ.get("AZURE_REAL_TEST", "false").lower() == "true"


def get_azure_clients():
    """Get Azure clients."""
    from azure.identity import ClientSecretCredential
    from azure.mgmt.compute import ComputeManagementClient
    from azure.mgmt.network import NetworkManagementClient
    from azure.mgmt.resource import ResourceManagementClient

    tenant_id = os.environ.get("AZURE_TENANT_ID")
    client_id = os.environ.get("AZURE_CLIENT_ID")
    client_secret = os.environ.get("AZURE_CLIENT_SECRET")
    subscription_id = os.environ.get("AZURE_SUBSCRIPTION_ID")

    credential = ClientSecretCredential(
        tenant_id=tenant_id,
        client_id=client_id,
        client_secret=client_secret,
    )

    return {
        "compute": ComputeManagementClient(credential, subscription_id),
        "network": NetworkManagementClient(credential, subscription_id),
        "resource": ResourceManagementClient(credential, subscription_id),
        "subscription_id": subscription_id,
    }


def ensure_resource_group(clients, resource_group: str, location: str):
    """Ensure the resource group exists."""
    resource_client = clients["resource"]
    try:
        resource_client.resource_groups.get(resource_group)
        print(f"   Resource group '{resource_group}' exists")
    except Exception:
        print(f"   Creating resource group '{resource_group}' in {location}...")
        resource_client.resource_groups.create_or_update(
            resource_group, {"location": location}
        )
        print(f"   Resource group created")


def ensure_virtual_network(clients, resource_group: str, location: str, vnet_name: str = None):
    """Ensure virtual network and subnet exist, return subnet ID."""
    network_client = clients["network"]
    subscription_id = clients["subscription_id"]
    # Use location-specific names to avoid cross-region issues
    if vnet_name is None:
        vnet_name = f"gateway-vnet-{location}"
    subnet_name = f"gateway-subnet-{location}"

    try:
        network_client.virtual_networks.get(resource_group, vnet_name)
        print(f"   VNet '{vnet_name}' exists")
    except Exception:
        print(f"   Creating virtual network '{vnet_name}'...")
        vnet_params = {
            "location": location,
            "address_space": {"address_prefixes": ["10.1.0.0/16"]},
            "subnets": [{"name": subnet_name, "address_prefix": "10.1.0.0/24"}],
        }
        poller = network_client.virtual_networks.begin_create_or_update(
            resource_group, vnet_name, vnet_params
        )
        poller.result()
        print(f"   VNet created")

    return f"/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.Network/virtualNetworks/{vnet_name}/subnets/{subnet_name}"


def create_public_ip(clients, resource_group: str, location: str, ip_name: str) -> str:
    """Create a public IP address and return its ID."""
    network_client = clients["network"]

    ip_params = {
        "location": location,
        "sku": {"name": "Standard"},
        "public_ip_allocation_method": "Static",
        "public_ip_address_version": "IPV4",
    }

    print(f"   Creating public IP '{ip_name}'...")
    poller = network_client.public_ip_addresses.begin_create_or_update(
        resource_group, ip_name, ip_params
    )
    ip_result = poller.result()
    print(f"   Public IP created: {ip_result.ip_address}")
    return ip_result.id


def create_network_interface(clients, resource_group: str, location: str, nic_name: str, subnet_id: str, public_ip_id: str) -> str:
    """Create a network interface and return its ID."""
    network_client = clients["network"]

    nic_params = {
        "location": location,
        "ip_configurations": [
            {
                "name": f"{nic_name}-ipconfig",
                "subnet": {"id": subnet_id},
                "public_ip_address": {"id": public_ip_id},
            }
        ],
    }

    print(f"   Creating network interface '{nic_name}'...")
    poller = network_client.network_interfaces.begin_create_or_update(
        resource_group, nic_name, nic_params
    )
    nic_result = poller.result()
    return nic_result.id


def create_vm(
    clients,
    resource_group: str,
    location: str,
    vm_name: str,
    vm_size: str,
    nic_id: str,
    startup_script: str = None,
):
    """Create a virtual machine."""
    compute_client = clients["compute"]

    # VM parameters
    vm_params = {
        "location": location,
        "tags": {
            "created-by": "gateway",
            "managed": "true",
            "purpose": "test",
        },
        "hardware_profile": {"vm_size": vm_size},
        "storage_profile": {
            "image_reference": {
                "publisher": "Canonical",
                "offer": "0001-com-ubuntu-server-jammy",
                "sku": "22_04-lts-gen2",
                "version": "latest",
            },
            "os_disk": {
                "name": f"{vm_name}-osdisk",
                "caching": "ReadWrite",
                "create_option": "FromImage",
                "managed_disk": {"storage_account_type": "Standard_LRS"},
            },
        },
        "os_profile": {
            "computer_name": vm_name,
            "admin_username": "azureuser",
            "admin_password": f"Gateway{int(time.time())}!",  # Temp password, will be rotated
            "linux_configuration": {
                "disable_password_authentication": False,
            },
        },
        "network_profile": {
            "network_interfaces": [{"id": nic_id}]
        },
    }

    # Add startup script via custom data (cloud-init)
    if startup_script:
        import base64
        # Wrap in cloud-init format
        cloud_init = f"""#cloud-config
runcmd:
  - |
    {startup_script}
"""
        vm_params["os_profile"]["custom_data"] = base64.b64encode(cloud_init.encode()).decode()

    print(f"   Creating VM '{vm_name}' ({vm_size})...")
    poller = compute_client.virtual_machines.begin_create_or_update(
        resource_group, vm_name, vm_params
    )
    vm_result = poller.result()
    print(f"   VM created: {vm_result.vm_id}")
    return vm_result


def get_vm_public_ip(clients, resource_group: str, ip_name: str) -> str:
    """Get the public IP address value."""
    network_client = clients["network"]
    ip = network_client.public_ip_addresses.get(resource_group, ip_name)
    return ip.ip_address


def get_boot_diagnostics(clients, resource_group: str, vm_name: str) -> str:
    """Get boot diagnostics / serial console output."""
    compute_client = clients["compute"]
    try:
        # This requires boot diagnostics to be enabled
        result = compute_client.virtual_machines.retrieve_boot_diagnostics_data(
            resource_group, vm_name
        )
        return str(result)
    except Exception as e:
        return f"Could not get boot diagnostics: {e}"


def delete_vm_and_resources(clients, resource_group: str, vm_name: str):
    """Delete VM and associated resources."""
    compute_client = clients["compute"]
    network_client = clients["network"]

    ip_name = f"{vm_name}-ip"
    nic_name = f"{vm_name}-nic"
    disk_name = f"{vm_name}-osdisk"

    # Delete VM
    print(f"   Deleting VM '{vm_name}'...")
    try:
        poller = compute_client.virtual_machines.begin_delete(resource_group, vm_name)
        poller.result()
        print(f"   VM deleted")
    except Exception as e:
        print(f"   Could not delete VM: {e}")

    # Wait a moment for VM deletion to complete
    time.sleep(5)

    # Delete NIC
    print(f"   Deleting network interface '{nic_name}'...")
    try:
        poller = network_client.network_interfaces.begin_delete(resource_group, nic_name)
        poller.result()
        print(f"   NIC deleted")
    except Exception as e:
        print(f"   Could not delete NIC: {e}")

    # Delete Public IP
    print(f"   Deleting public IP '{ip_name}'...")
    try:
        poller = network_client.public_ip_addresses.begin_delete(resource_group, ip_name)
        poller.result()
        print(f"   Public IP deleted")
    except Exception as e:
        print(f"   Could not delete public IP: {e}")

    # Delete OS disk
    print(f"   Deleting OS disk '{disk_name}'...")
    try:
        poller = compute_client.disks.begin_delete(resource_group, disk_name)
        poller.result()
        print(f"   OS disk deleted")
    except Exception as e:
        print(f"   Could not delete OS disk: {e}")


def main():
    """Launch a test VM, run compute, and clean up."""

    # Safety check - don't run real tests by default
    if not REAL_TEST:
        print("=" * 60)
        print("Azure VM Launch Test - SKIPPED (REAL_TEST=False)")
        print("=" * 60)
        print("  This test creates REAL cloud resources that cost money!")
        print("  To run this test, set the environment variable:")
        print("    export AZURE_REAL_TEST=true")
        print("  Or modify REAL_TEST in this file.")
        print("=" * 60)
        print("\nTest skipped successfully (no resources created)")
        return 0

    resource_group = os.environ.get("AZURE_RESOURCE_GROUP", "gateway-test-rg")
    location = os.environ.get("AZURE_LOCATION", "eastus")
    vm_name = f"gateway-test-{int(time.time())}"
    vm_size = "Standard_B1ms"

    print("=" * 60)
    print("GateWay Azure VM Launch Test")
    print("=" * 60)
    print(f"  Resource Group: {resource_group}")
    print(f"  Location: {location}")
    print(f"  VM Name: {vm_name}")
    print(f"  VM Size: {vm_size} (FREE TIER)")
    print("  REAL_TEST=True - This WILL create cloud resources!")
    print("=" * 60)

    # Startup script that does a simple compute job
    startup_script = '''#!/bin/bash
echo "=========================================="
echo "GateWay Test Job Starting"
echo "=========================================="
echo "Hostname: $(hostname)"
echo "Date: $(date)"
echo "CPU Info: $(nproc) cores"
echo "Memory: $(free -h | grep Mem | awk '{print $2}')"
echo ""
echo "Running compute test..."

# Simple Python compute test
python3 -c "
import math
import time

print('Computing sum of square roots from 0 to 99999...')
start = time.time()
result = sum(math.sqrt(i) for i in range(100000))
elapsed = time.time() - start
print(f'Result: {result:.2f}')
print(f'Time: {elapsed:.3f} seconds')

print('')
print('Computing prime numbers up to 10000...')
start = time.time()
primes = []
for n in range(2, 10001):
    is_prime = True
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0:
            is_prime = False
            break
    if is_prime:
        primes.append(n)
elapsed = time.time() - start
print(f'Found {len(primes)} primes')
print(f'Last 5 primes: {primes[-5:]}')
print(f'Time: {elapsed:.3f} seconds')
"

echo ""
echo "=========================================="
echo "GateWay Test Job Complete!"
echo "=========================================="

# Signal completion
touch /tmp/gateway_job_complete
'''

    try:
        clients = get_azure_clients()
    except ImportError as e:
        print(f"\nMissing Azure SDK: {e}")
        print("Install with: pip install azure-identity azure-mgmt-compute azure-mgmt-network azure-mgmt-resource")
        return 1

    vm_created = False

    try:
        # Setup infrastructure
        print("\nSetting up infrastructure...")
        ensure_resource_group(clients, resource_group, location)
        subnet_id = ensure_virtual_network(clients, resource_group, location)

        # Create networking resources
        print("\nCreating networking resources...")
        ip_name = f"{vm_name}-ip"
        nic_name = f"{vm_name}-nic"

        public_ip_id = create_public_ip(clients, resource_group, location, ip_name)
        nic_id = create_network_interface(clients, resource_group, location, nic_name, subnet_id, public_ip_id)

        # Create VM
        print("\nCreating VM...")
        create_vm(
            clients=clients,
            resource_group=resource_group,
            location=location,
            vm_name=vm_name,
            vm_size=vm_size,
            nic_id=nic_id,
            startup_script=startup_script,
        )
        vm_created = True

        # Get public IP
        public_ip = get_vm_public_ip(clients, resource_group, ip_name)
        print(f"   Public IP: {public_ip}")

        # Wait for job to complete
        print("\nWaiting for compute job to complete...")
        print("   (Azure VMs take longer to provision, waiting up to 5 minutes)")
        print("   Note: Unlike GCP, Azure doesn't easily expose serial console output")
        print("   The job should be running via cloud-init...")

        # Since Azure doesn't have easy serial console access like GCP,
        # we'll just wait a reasonable time for the job to complete
        max_wait = 300  # 5 minutes
        waited = 0
        wait_interval = 30

        while waited < max_wait:
            time.sleep(wait_interval)
            waited += wait_interval
            print(f"   Waiting... ({waited}s / {max_wait}s)")

            # Try to get boot diagnostics (may not always work)
            diagnostics = get_boot_diagnostics(clients, resource_group, vm_name)
            if "GateWay Test Job Complete" in diagnostics:
                print("\n" + "=" * 60)
                print("JOB OUTPUT FROM CLOUD VM:")
                print("=" * 60)
                print(diagnostics)
                print("=" * 60)
                break

        if waited >= max_wait:
            print(f"\nReached wait timeout ({max_wait}s)")
            print("   The job should have completed by now")
            print("   You can SSH to the VM to verify: ssh azureuser@" + public_ip)

    finally:
        # Always clean up — even if VM was never created, network resources may exist
        print("\nCleaning up...")
        if vm_created:
            try:
                delete_vm_and_resources(clients, resource_group, vm_name)
            except Exception as e:
                print(f"   Cleanup failed: {e}")
                print(f"   Manual cleanup:")
                print(f"     az vm delete --resource-group {resource_group} --name {vm_name} --yes")
                print(f"     az network nic delete --resource-group {resource_group} --name {vm_name}-nic")
                print(f"     az network public-ip delete --resource-group {resource_group} --name {vm_name}-ip")
                print(f"     az disk delete --resource-group {resource_group} --name {vm_name}-osdisk --yes")
        else:
            # VM was never created, but NIC/IP may have been — clean them up
            print("   VM was not created, cleaning up any orphaned network resources...")
            network_client = clients["network"]
            for resource_name, delete_fn in [
                (f"{vm_name}-nic", lambda n: network_client.network_interfaces.begin_delete(resource_group, n)),
                (f"{vm_name}-ip", lambda n: network_client.public_ip_addresses.begin_delete(resource_group, n)),
            ]:
                try:
                    poller = delete_fn(resource_name)
                    poller.result()
                    print(f"   Deleted {resource_name}")
                except Exception:
                    pass  # Resource may not exist, that's fine

    print("\n" + "=" * 60)
    print("TEST COMPLETE!")
    print("=" * 60)
    print("Successfully launched VM on Azure")
    print("Ran compute job in the cloud (via cloud-init)")
    print("Cleaned up resources")
    print("\nGateWay is ready for Azure cloud execution!")

    return 0


if __name__ == "__main__":
    sys.exit(main())
