#!/usr/bin/env python3
"""
Azure Connection Test Script
Tests if GateWay can connect to Azure and execute simple compute tasks.

Note: This test makes API calls to Azure (read-only operations like listing VM sizes).
Set AZURE_REAL_TEST=true to enable API calls.
"""

import os
import sys

# =============================================================================
# SAFETY FLAG - Set to True to make actual Azure API calls
# =============================================================================
REAL_TEST = os.environ.get("AZURE_REAL_TEST", "false").lower() == "true"


def check_env():
    """Check Azure environment variables."""
    print("=" * 60)
    print("CHECKING AZURE ENVIRONMENT")
    print("=" * 60)

    subscription_id = os.environ.get("AZURE_SUBSCRIPTION_ID")
    tenant_id = os.environ.get("AZURE_TENANT_ID")
    client_id = os.environ.get("AZURE_CLIENT_ID")
    client_secret = os.environ.get("AZURE_CLIENT_SECRET")
    resource_group = os.environ.get("AZURE_RESOURCE_GROUP")
    location = os.environ.get("AZURE_LOCATION", "eastus")

    print(f"  AZURE_SUBSCRIPTION_ID: {'***' + subscription_id[-4:] if subscription_id else 'NOT SET'}")
    print(f"  AZURE_TENANT_ID: {'***' + tenant_id[-4:] if tenant_id else 'NOT SET'}")
    print(f"  AZURE_CLIENT_ID: {'***' + client_id[-4:] if client_id else 'NOT SET'}")
    print(f"  AZURE_CLIENT_SECRET: {'***' + client_secret[-4:] if client_secret else 'NOT SET'}")
    print(f"  AZURE_RESOURCE_GROUP: {resource_group or 'NOT SET (will default to gateway-rg)'}")
    print(f"  AZURE_LOCATION: {location}")

    if subscription_id and tenant_id and client_id and client_secret:
        print("  Credentials configured")
        return True
    else:
        print("  Credentials not configured")
        return False


def test_azure_auth():
    """Test Azure authentication."""
    print("\n" + "=" * 60)
    print("TESTING AZURE AUTHENTICATION")
    print("=" * 60)

    try:
        from azure.identity import ClientSecretCredential
        from azure.mgmt.resource import SubscriptionClient

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

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

        # Get subscription info to verify credentials
        subscription_client = SubscriptionClient(credential)
        subscription_id = os.environ.get("AZURE_SUBSCRIPTION_ID")

        subscription = subscription_client.subscriptions.get(subscription_id)
        print(f"  Subscription: {subscription.display_name}")
        print(f"  Subscription ID: {subscription.subscription_id}")
        print(f"  State: {subscription.state}")
        print("  Authentication successful!")
        return True
    except ImportError as e:
        print(f"  Authentication failed: Missing Azure SDK - {e}")
        print("  Install with: pip install azure-identity azure-mgmt-resource azure-mgmt-compute azure-mgmt-network")
        return False
    except Exception as e:
        print(f"  Authentication failed: {e}")
        return False


def test_compute_api():
    """Test Azure Compute API access."""
    print("\n" + "=" * 60)
    print("TESTING AZURE COMPUTE API")
    print("=" * 60)

    location = os.environ.get("AZURE_LOCATION", "eastus")

    try:
        from azure.identity import ClientSecretCredential
        from azure.mgmt.compute import ComputeManagementClient

        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,
        )

        compute_client = ComputeManagementClient(credential, subscription_id)

        # Test 1: List VM sizes
        print(f"\n  Listing available VM sizes in {location}...")
        vm_sizes = list(compute_client.virtual_machine_sizes.list(location))

        # Filter to common sizes
        common_sizes = []
        target_names = ["Standard_B1s", "Standard_B1ms", "Standard_B2s", "Standard_D2s_v3"]
        for size in vm_sizes:
            if size.name in target_names:
                common_sizes.append({
                    "name": size.name,
                    "vcpus": size.number_of_cores,
                    "memory_mb": size.memory_in_mb,
                })

        print(f"  Found {len(vm_sizes)} VM sizes total")

        # Show free tier option
        free_tier = [s for s in common_sizes if s["name"] == "Standard_B1s"]
        if free_tier:
            s = free_tier[0]
            print("\n  FREE TIER AVAILABLE:")
            memory_gb = s['memory_mb'] / 1024
            print(f"     {s['name']}: {s['vcpus']} vCPUs, {memory_gb:.1f}GB RAM")

        # Show sample sizes
        print("\n  Sample VM sizes:")
        for s in common_sizes[:5]:
            memory_gb = s['memory_mb'] / 1024
            print(f"     {s['name']}: {s['vcpus']} vCPUs, {memory_gb:.1f}GB RAM")

        # Test 2: List current VMs with gateway tag
        print(f"\n  Checking existing gateway VMs...")
        resource_group = os.environ.get("AZURE_RESOURCE_GROUP", "gateway-rg")

        try:
            vms = list(compute_client.virtual_machines.list(resource_group))
            gateway_vms = []
            for vm in vms:
                if vm.tags and vm.tags.get("created-by") == "gateway":
                    gateway_vms.append(vm)

            print(f"  Currently {len(gateway_vms)} gateway VMs in {resource_group}")
            for vm in gateway_vms:
                print(f"     - {vm.name}: {vm.provisioning_state}")
        except Exception as e:
            # Resource group may not exist yet
            if "ResourceGroupNotFound" in str(e):
                print(f"  Resource group '{resource_group}' does not exist (will be created when needed)")
            else:
                print(f"  Could not list VMs: {e}")

        return True

    except ImportError as e:
        print(f"  Compute API test failed: Missing Azure SDK - {e}")
        return False
    except Exception as e:
        print(f"  Compute API test failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_simple_compute():
    """Simulate what a GateWay job would do."""
    print("\n" + "=" * 60)
    print("SIMULATING GATEWAY JOB (Local Compute)")
    print("=" * 60)

    # This simulates what would run on the cloud instance
    print("\n  Running simple compute task locally...")

    # Simple compute test
    import math
    result = sum(math.sqrt(i) for i in range(100000))
    print(f"  Computed: sum(sqrt(0..99999)) = {result:.2f}")

    # Test serialization (what GateWay uses)
    try:
        import pickle

        import cloudpickle

        def sample_task(data):
            """Simulated compute task."""
            return sum(x ** 2 for x in data)

        serialized = cloudpickle.dumps(sample_task)
        restored = pickle.loads(serialized)
        test_result = restored([1, 2, 3, 4, 5])

        print(f"  Cloudpickle serialization works: {test_result}")
    except ImportError:
        print("  cloudpickle not installed (pip install cloudpickle)")

    return True


def main():
    """Run all tests."""
    print("\n" + "GateWay Azure Connection Test")
    print("=" * 60 + "\n")

    # Safety check - don't make real API calls by default
    if not REAL_TEST:
        print("Azure Connection Test - SKIPPED (REAL_TEST=False)")
        print("=" * 60)
        print("  This test makes API calls to Azure.")
        print("  To run this test, set the environment variable:")
        print("    export AZURE_REAL_TEST=true")
        print("=" * 60)
        print("\nTest skipped successfully (no API calls made)")
        return 0

    results = {}

    # Check environment
    results["env"] = check_env()

    if not results["env"]:
        print("\nEnvironment not configured. Please set:")
        print("   export AZURE_SUBSCRIPTION_ID=your-subscription-id")
        print("   export AZURE_TENANT_ID=your-tenant-id")
        print("   export AZURE_CLIENT_ID=your-client-id")
        print("   export AZURE_CLIENT_SECRET=your-client-secret")
        print("   export AZURE_RESOURCE_GROUP=gateway-rg")
        print("   export AZURE_LOCATION=eastus")
        return 1

    # Test auth
    results["auth"] = test_azure_auth()

    # Test Compute API
    results["compute_api"] = test_compute_api()

    # Test local compute simulation
    results["local_compute"] = test_simple_compute()

    # Summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)

    all_passed = True
    for test, passed in results.items():
        status = "PASS" if passed else "FAIL"
        print(f"  {test}: {status}")
        if not passed:
            all_passed = False

    if all_passed:
        print("\nALL TESTS PASSED!")
        print("\nAzure is ready for GateWay!")
        print("   Your free tier setup can launch Standard_B1s instances")
        print("   for $0.00/hour (750 hours/month for 12 months)")
    else:
        print("\nSome tests failed. Check the errors above.")

    return 0 if all_passed else 1


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