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

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

import json
import os
import sys
import time

from google.cloud import compute_v1
from google.oauth2 import service_account

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


def get_credentials():
    """Load GCP credentials from file path or JSON string."""
    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 create_instance(
    project_id: str,
    zone: str,
    instance_name: str,
    machine_type: str = "e2-micro",
    startup_script: str = None,
) -> compute_v1.Instance:
    """Create a new GCE instance."""

    credentials = get_credentials()
    instances_client = compute_v1.InstancesClient(credentials=credentials)

    # Use Debian 12 (bookworm) - free tier compatible
    image_response = compute_v1.ImagesClient(credentials=credentials).get_from_family(
        project="debian-cloud",
        family="debian-12"
    )
    source_disk_image = image_response.self_link

    # Configure the machine
    machine_type_path = f"zones/{zone}/machineTypes/{machine_type}"

    # Boot disk
    disk = compute_v1.AttachedDisk()
    disk.boot = True
    disk.auto_delete = True
    disk.initialize_params = compute_v1.AttachedDiskInitializeParams()
    disk.initialize_params.source_image = source_disk_image
    disk.initialize_params.disk_size_gb = 10  # Minimal disk for free tier
    disk.initialize_params.disk_type = f"zones/{zone}/diskTypes/pd-standard"

    # Network interface
    network_interface = compute_v1.NetworkInterface()
    network_interface.network = "global/networks/default"
    access_config = compute_v1.AccessConfig()
    access_config.name = "External NAT"
    access_config.type_ = "ONE_TO_ONE_NAT"
    network_interface.access_configs = [access_config]

    # Instance configuration
    instance = compute_v1.Instance()
    instance.name = instance_name
    instance.machine_type = machine_type_path
    instance.disks = [disk]
    instance.network_interfaces = [network_interface]

    # Labels for identification (must include "managed=true" for orphan cleanup)
    instance.labels = {
        "created-by": "gateway",
        "managed": "true",
        "purpose": "test",
    }

    # Startup script to run compute job
    if startup_script:
        metadata = compute_v1.Metadata()
        metadata.items = [
            compute_v1.Items(key="startup-script", value=startup_script)
        ]
        instance.metadata = metadata

    # Create the instance
    print(f"📦 Creating instance '{instance_name}' in {zone}...")
    operation = instances_client.insert(
        project=project_id,
        zone=zone,
        instance_resource=instance,
    )

    return operation, instances_client


def wait_for_operation(project_id: str, zone: str, operation_name: str, credentials):
    """Wait for a zone operation to complete."""
    operations_client = compute_v1.ZoneOperationsClient(credentials=credentials)

    while True:
        result = operations_client.get(
            project=project_id,
            zone=zone,
            operation=operation_name,
        )

        if result.status == compute_v1.Operation.Status.DONE:
            if result.error:
                raise Exception(f"Operation failed: {result.error}")
            return result

        print("   ⏳ Waiting for instance to be ready...")
        time.sleep(3)


def get_instance_ip(project_id: str, zone: str, instance_name: str, credentials) -> str:
    """Get the external IP of an instance."""
    instances_client = compute_v1.InstancesClient(credentials=credentials)
    instance = instances_client.get(
        project=project_id,
        zone=zone,
        instance=instance_name,
    )

    for interface in instance.network_interfaces:
        for config in interface.access_configs:
            if config.nat_i_p:
                return config.nat_i_p
    return None


def delete_instance(project_id: str, zone: str, instance_name: str, credentials):
    """Delete a GCE instance."""
    instances_client = compute_v1.InstancesClient(credentials=credentials)

    print(f"🗑️  Deleting instance '{instance_name}'...")
    operation = instances_client.delete(
        project=project_id,
        zone=zone,
        instance=instance_name,
    )

    wait_for_operation(project_id, zone, operation.name, credentials)
    print("   ✅ Instance deleted")


def get_serial_output(project_id: str, zone: str, instance_name: str, credentials) -> str:
    """Get serial port output from instance."""
    instances_client = compute_v1.InstancesClient(credentials=credentials)

    try:
        output = instances_client.get_serial_port_output(
            project=project_id,
            zone=zone,
            instance=instance_name,
        )
        return output.contents
    except Exception as e:
        return f"Could not get serial output: {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("⚠️  GCP 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 GCP_REAL_TEST=true")
        print("  Or modify REAL_TEST in this file.")
        print("=" * 60)
        print("\n✅ Test skipped successfully (no resources created)")
        return 0

    project_id = os.environ.get("GCP_PROJECT_ID")
    zone = os.environ.get("GCP_ZONE", "us-central1-a")
    instance_name = f"gateway-test-{int(time.time())}"

    print("=" * 60)
    print("🚀 GateWay GCP VM Launch Test")
    print("=" * 60)
    print(f"  Project: {project_id}")
    print(f"  Zone: {zone}")
    print(f"  Instance: {instance_name}")
    print("  Type: e2-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
'''

    credentials = get_credentials()

    try:
        # Create the instance
        operation, instances_client = create_instance(
            project_id=project_id,
            zone=zone,
            instance_name=instance_name,
            machine_type="e2-micro",
            startup_script=startup_script,
        )

        # Wait for instance to be created
        print("   ⏳ Waiting for instance creation...")
        wait_for_operation(project_id, zone, operation.name, credentials)
        print("   ✅ Instance created!")

        # Get external IP
        ip = get_instance_ip(project_id, zone, instance_name, credentials)
        print(f"   🌐 External IP: {ip}")

        # Wait for startup script to complete (check serial output)
        print("\n📋 Waiting for compute job to complete...")
        print("   (Checking serial console output every 10 seconds)")

        job_complete = False
        max_wait = 120  # 2 minutes max
        waited = 0

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

            serial_output = get_serial_output(project_id, zone, instance_name, credentials)

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

                # Extract just the job output
                lines = serial_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("⚠️  Job didn't complete in time, but that's OK for a test")
            print("   Serial output preview:")
            print(serial_output[-1000:] if len(serial_output) > 1000 else serial_output)

    finally:
        # Always clean up
        print("\n🧹 Cleaning up...")
        try:
            delete_instance(project_id, zone, instance_name, credentials)
        except Exception as e:
            print(f"   ⚠️  Cleanup failed: {e}")
            print(f"   Manual cleanup: gcloud compute instances delete {instance_name} --zone={zone}")

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

    return 0


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