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

Note: This test makes API calls to GCP (read-only operations like listing machine types).
Set REAL_TEST = True to enable API calls.
"""

import os
import sys

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


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

    project_id = os.environ.get("GCP_PROJECT_ID")
    zone = os.environ.get("GCP_ZONE")
    creds_value = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")

    print(f"  GCP_PROJECT_ID: {project_id or 'NOT SET'}")
    print(f"  GCP_ZONE: {zone or 'NOT SET'}")

    if creds_value:
        # Check if it's a JSON string or file path
        if creds_value.strip().startswith("{"):
            print("  GOOGLE_APPLICATION_CREDENTIALS: JSON string (inline)")
            print("  Credentials configured (JSON)")
        elif os.path.exists(creds_value):
            print(f"  GOOGLE_APPLICATION_CREDENTIALS: {creds_value}")
            print("  Credentials file exists")
        else:
            print(f"  GOOGLE_APPLICATION_CREDENTIALS: {creds_value}")
            print(f"  Credentials file NOT FOUND")
            return False
    else:
        print("  GOOGLE_APPLICATION_CREDENTIALS: NOT SET")
        return False

    return bool(project_id and zone and creds_value)


def get_gcp_credentials():
    """Get GCP credentials from file path or JSON string."""
    import json
    from google.oauth2 import service_account

    creds_value = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
    if not creds_value:
        raise ValueError("GOOGLE_APPLICATION_CREDENTIALS not set")

    # Check if it's a JSON string or file path
    if creds_value.strip().startswith("{"):
        # It's a JSON string
        creds_info = json.loads(creds_value)
        return service_account.Credentials.from_service_account_info(
            creds_info,
            scopes=["https://www.googleapis.com/auth/cloud-platform"]
        )
    else:
        # It's a file path
        return service_account.Credentials.from_service_account_file(
            creds_value,
            scopes=["https://www.googleapis.com/auth/cloud-platform"]
        )


def test_gcp_auth():
    """Test GCP authentication."""
    print("\n" + "=" * 60)
    print("TESTING GCP AUTHENTICATION")
    print("=" * 60)

    try:
        credentials = get_gcp_credentials()
        project_id = os.environ.get("GCP_PROJECT_ID")
        print("  Authentication successful!")
        print(f"  Project: {project_id}")
        print(f"  Service Account: {credentials.service_account_email}")
        return True
    except Exception as e:
        print(f"  Authentication failed: {e}")
        return False


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

    project_id = os.environ.get("GCP_PROJECT_ID")
    zone = os.environ.get("GCP_ZONE")

    try:
        from google.cloud import compute_v1

        credentials = get_gcp_credentials()

        # Test 1: List machine types (doesn't require any instances)
        print("\n  Listing available machine types...")
        machine_types_client = compute_v1.MachineTypesClient(credentials=credentials)
        request = compute_v1.ListMachineTypesRequest(
            project=project_id,
            zone=zone,
        )

        machine_types = []
        for mt in machine_types_client.list(request=request):
            machine_types.append({
                "name": mt.name,
                "vcpus": mt.guest_cpus,
                "memory_mb": mt.memory_mb,
            })

        print(f"  Found {len(machine_types)} machine types in {zone}")

        # Show free tier option
        free_tier = [m for m in machine_types if m["name"] == "e2-micro"]
        if free_tier:
            mt = free_tier[0]
            print("\n  FREE TIER AVAILABLE:")
            print(f"     e2-micro: {mt['vcpus']} vCPUs, {mt['memory_mb']}MB RAM")

        # Show a few options
        print("\n  Sample machine types:")
        for mt in machine_types[:5]:
            memory_gb = mt['memory_mb'] / 1024
            print(f"     {mt['name']}: {mt['vcpus']} vCPUs, {memory_gb:.1f}GB RAM")

        # Test 2: List current instances
        print("\n  Checking existing instances...")
        instances_client = compute_v1.InstancesClient(credentials=credentials)
        request = compute_v1.ListInstancesRequest(
            project=project_id,
            zone=zone,
        )

        instances = list(instances_client.list(request=request))
        print(f"  Currently {len(instances)} instances in {zone}")
        for inst in instances:
            print(f"     - {inst.name}: {inst.status}")

        return True

    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_gpu_task(data):
            """Simulated GPU task."""
            return sum(x ** 2 for x in data)

        serialized = cloudpickle.dumps(sample_gpu_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 GCP Connection Test")
    print("=" * 60 + "\n")

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

    results = {}

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

    if not results["env"]:
        print("\n❌ Environment not configured. Please set:")
        print("   export GCP_PROJECT_ID=your-project-id")
        print("   export GCP_ZONE=us-central1-a")
        print("   export GOOGLE_APPLICATION_CREDENTIALS=/path/to/key.json")
        return 1

    # Test auth
    results["auth"] = test_gcp_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("\n🎉 ALL TESTS PASSED!")
        print("\n✅ GCP is ready for GateWay!")
        print("   Your free tier setup can launch e2-micro instances")
        print("   in us-central1-a for $0.00/hour")
    else:
        print("\n⚠️  Some tests failed. Check the errors above.")

    return 0 if all_passed else 1


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