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

WARNING: This test creates REAL cloud resources that cost money!
Set AWS_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("AWS_REAL_TEST", "false").lower() == "true"


def get_ec2_client(region: str):
    """Get boto3 EC2 client."""
    import boto3
    return boto3.client("ec2", region_name=region)


def get_default_ami(ec2_client, region: str) -> str:
    """Get the latest Amazon Linux 2023 AMI for the region."""
    response = ec2_client.describe_images(
        Owners=["amazon"],
        Filters=[
            {"Name": "name", "Values": ["al2023-ami-2023*-x86_64"]},
            {"Name": "state", "Values": ["available"]},
            {"Name": "architecture", "Values": ["x86_64"]},
        ],
    )

    # Sort by creation date and get the latest
    images = sorted(response["Images"], key=lambda x: x["CreationDate"], reverse=True)
    if images:
        return images[0]["ImageId"]

    # Fallback to known AMIs
    fallback_amis = {
        "us-east-1": "ami-0c7217cdde317cfec",
        "us-east-2": "ami-0900fe555666598a2",
        "us-west-1": "ami-0f5ee92b1cea87159",
        "us-west-2": "ami-008fe2fc65df48dac",
    }
    return fallback_amis.get(region, fallback_amis["us-east-1"])


def create_instance(
    ec2_client,
    instance_name: str,
    instance_type: str = "t3.micro",
    ami_id: str = None,
    startup_script: str = None,
) -> str:
    """Create a new EC2 instance."""

    # Build instance parameters
    run_params = {
        "ImageId": ami_id,
        "InstanceType": instance_type,
        "MinCount": 1,
        "MaxCount": 1,
        "TagSpecifications": [
            {
                "ResourceType": "instance",
                "Tags": [
                    {"Key": "Name", "Value": instance_name},
                    {"Key": "created-by", "Value": "gateway"},
                    {"Key": "managed", "Value": "true"},
                    {"Key": "purpose", "Value": "test"},
                ],
            }
        ],
    }

    if startup_script:
        run_params["UserData"] = startup_script

    print(f"Creating instance '{instance_name}'...")
    response = ec2_client.run_instances(**run_params)
    instance_id = response["Instances"][0]["InstanceId"]

    return instance_id


def wait_for_instance_running(ec2_client, instance_id: str, max_wait: int = 120):
    """Wait for instance to be running."""
    print("   Waiting for instance to be running...")
    waiter = ec2_client.get_waiter("instance_running")
    waiter.wait(
        InstanceIds=[instance_id],
        WaiterConfig={"Delay": 5, "MaxAttempts": max_wait // 5}
    )
    print("   Instance is running")


def get_instance_ip(ec2_client, instance_id: str) -> str:
    """Get the public IP of an instance."""
    response = ec2_client.describe_instances(InstanceIds=[instance_id])
    instance = response["Reservations"][0]["Instances"][0]
    return instance.get("PublicIpAddress")


def get_console_output(ec2_client, instance_id: str) -> str:
    """Get console output from instance."""
    try:
        response = ec2_client.get_console_output(InstanceId=instance_id)
        output = response.get("Output", "")
        if output:
            import base64
            try:
                return base64.b64decode(output).decode("utf-8", errors="ignore")
            except Exception:
                return output
        return ""
    except Exception as e:
        return f"Could not get console output: {e}"


def terminate_instance(ec2_client, instance_id: str):
    """Terminate an EC2 instance."""
    print(f"Deleting instance '{instance_id}'...")
    ec2_client.terminate_instances(InstanceIds=[instance_id])

    # Wait for termination
    waiter = ec2_client.get_waiter("instance_terminated")
    waiter.wait(InstanceIds=[instance_id], WaiterConfig={"Delay": 5, "MaxAttempts": 24})
    print("   Instance terminated")


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("AWS 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 AWS_REAL_TEST=true")
        print("  Or modify REAL_TEST in this file.")
        print("=" * 60)
        print("\nTest skipped successfully (no resources created)")
        return 0

    region = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
    instance_name = f"gateway-test-{int(time.time())}"

    print("=" * 60)
    print("GateWay AWS VM Launch Test")
    print("=" * 60)
    print(f"  Region: {region}")
    print(f"  Instance: {instance_name}")
    print("  Type: t3.micro (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
'''

    ec2_client = get_ec2_client(region)
    instance_id = None

    try:
        # Get AMI
        print("\n  Getting latest Amazon Linux 2023 AMI...")
        ami_id = get_default_ami(ec2_client, region)
        print(f"  Using AMI: {ami_id}")

        # Create the instance
        instance_id = create_instance(
            ec2_client=ec2_client,
            instance_name=instance_name,
            instance_type="t3.micro",
            ami_id=ami_id,
            startup_script=startup_script,
        )
        print(f"   Instance ID: {instance_id}")

        # Wait for instance to be running
        wait_for_instance_running(ec2_client, instance_id)

        # Get external IP
        ip = get_instance_ip(ec2_client, instance_id)
        print(f"   Public IP: {ip}")

        # Wait for startup script to complete (check console output)
        print("\nWaiting for compute job to complete...")
        print("   (Checking console output every 15 seconds)")

        job_complete = False
        max_wait = 180  # 3 minutes max
        waited = 0

        while not job_complete and waited < max_wait:
            time.sleep(15)
            waited += 15

            console_output = get_console_output(ec2_client, instance_id)

            if "GateWay Test Job Complete" in console_output:
                job_complete = True
                print("\n" + "=" * 60)
                print("JOB OUTPUT FROM CLOUD VM:")
                print("=" * 60)

                # Extract just the job output
                lines = console_output.split('\n')
                in_job = False
                for line in lines:
                    if "GateWay Test Job Starting" in line:
                        in_job = True
                    if in_job:
                        print(f"  {line}")
                    if "GateWay Test Job Complete" in line:
                        break

                print("=" * 60)
            else:
                print(f"   Still running... ({waited}s)")

        if not job_complete:
            print(f"Job didn't complete in time ({max_wait}s), but that's OK for a test")
            if console_output:
                print("   Console output preview:")
                print(console_output[-1000:] if len(console_output) > 1000 else console_output)

    finally:
        # Always clean up
        print("\nCleaning up...")
        if instance_id:
            try:
                terminate_instance(ec2_client, instance_id)
            except Exception as e:
                print(f"   Cleanup failed: {e}")
                print(f"   Manual cleanup: aws ec2 terminate-instances --instance-ids {instance_id}")

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

    return 0


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