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

Note: This test makes API calls to AWS (read-only operations like listing instance types).
Set AWS_REAL_TEST=true to enable API calls.
"""

import os
import sys

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


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

    access_key = os.environ.get("AWS_ACCESS_KEY_ID")
    secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
    region = os.environ.get("AWS_DEFAULT_REGION")

    print(f"  AWS_ACCESS_KEY_ID: {'***' + access_key[-4:] if access_key else 'NOT SET'}")
    print(f"  AWS_SECRET_ACCESS_KEY: {'***' + secret_key[-4:] if secret_key else 'NOT SET'}")
    print(f"  AWS_DEFAULT_REGION: {region or 'NOT SET (will default to us-east-1)'}")

    if access_key and secret_key:
        print("  Credentials configured")
        return True
    else:
        print("  Credentials not configured")
        return False


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

    try:
        import boto3
        from botocore.exceptions import ClientError, NoCredentialsError

        # Simple STS call to verify credentials
        sts = boto3.client("sts")
        identity = sts.get_caller_identity()
        print(f"  Account: {identity['Account']}")
        print(f"  User ARN: {identity['Arn']}")
        print("  Authentication successful!")
        return True
    except NoCredentialsError:
        print("  Authentication failed: No credentials found")
        return False
    except ClientError as e:
        print(f"  Authentication failed: {e}")
        return False
    except Exception as e:
        print(f"  Authentication failed: {e}")
        return False


def test_ec2_api():
    """Test EC2 API access."""
    print("\n" + "=" * 60)
    print("TESTING EC2 API")
    print("=" * 60)

    region = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")

    try:
        import boto3

        ec2 = boto3.client("ec2", region_name=region)

        # Test 1: List instance types
        print(f"\n  Listing available instance types in {region}...")
        response = ec2.describe_instance_types(
            Filters=[
                {"Name": "instance-type", "Values": ["t2.micro", "t3.micro", "t3.small", "t3.medium"]}
            ]
        )

        instance_types = []
        for it in response["InstanceTypes"]:
            instance_types.append({
                "name": it["InstanceType"],
                "vcpus": it["VCpuInfo"]["DefaultVCpus"],
                "memory_mb": it["MemoryInfo"]["SizeInMiB"],
            })

        print(f"  Found {len(instance_types)} instance types")

        # Show free tier option
        free_tier = [i for i in instance_types if i["name"] in ("t2.micro", "t3.micro")]
        if free_tier:
            print("\n  FREE TIER AVAILABLE:")
            for it in free_tier:
                memory_gb = it['memory_mb'] / 1024
                print(f"     {it['name']}: {it['vcpus']} vCPUs, {memory_gb:.1f}GB RAM")

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

        # Test 2: List current instances
        print(f"\n  Checking existing instances in {region}...")
        instances_response = ec2.describe_instances(
            Filters=[
                {"Name": "tag:created-by", "Values": ["gateway"]}
            ]
        )

        instance_count = 0
        for reservation in instances_response["Reservations"]:
            for inst in reservation["Instances"]:
                if inst["State"]["Name"] not in ("terminated", "shutting-down"):
                    instance_count += 1
                    print(f"     - {inst['InstanceId']}: {inst['State']['Name']}")

        print(f"  Currently {instance_count} gateway instances in {region}")

        return True

    except Exception as e:
        print(f"  EC2 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 AWS Connection Test")
    print("=" * 60 + "\n")

    # Safety check - don't make real API calls by default
    if not REAL_TEST:
        print("AWS Connection Test - SKIPPED (REAL_TEST=False)")
        print("=" * 60)
        print("  This test makes API calls to AWS.")
        print("  To run this test, set the environment variable:")
        print("    export AWS_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 AWS_ACCESS_KEY_ID=your-access-key")
        print("   export AWS_SECRET_ACCESS_KEY=your-secret-key")
        print("   export AWS_DEFAULT_REGION=us-east-1")
        return 1

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

    # Test EC2 API
    results["ec2_api"] = test_ec2_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("\nAWS is ready for GateWay!")
        print("   Your free tier setup can launch t2.micro/t3.micro instances")
        print("   in us-east-1 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())
