#!/usr/bin/env python3
"""
GateWay End-to-End Test

Tests the complete flow of submitting a job and executing it on GCP.
This uses the GCP cloud provider integration directly.

WARNING: These tests create REAL cloud resources that cost money!
Set GATEWAY_E2E_TESTS=true only when you explicitly want to run E2E tests.
"""

import asyncio
import os
import pickle
import sys
import time

import cloudpickle
import pytest

# Add parent to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from verlex_server.gateway.backend.cloud_providers import GCPProvider, get_provider

# =============================================================================
# SAFETY FLAG - Set to True only when you want to run real E2E tests
# =============================================================================
REAL_TEST = os.environ.get("GATEWAY_E2E_TESTS", "").lower() in ("true", "1", "yes")

# Skip all tests in this module unless explicitly enabled (for pytest)
pytestmark = pytest.mark.skipif(
    not REAL_TEST,
    reason="E2E tests require GCP infrastructure. Set GATEWAY_E2E_TESTS=true to run."
)


def print_section(title: str):
    """Print a section header."""
    print("\n" + "=" * 60)
    print(f"  {title}")
    print("=" * 60)


# ============================================================================
# Test Functions - These would normally be user-defined
# ============================================================================

def simple_compute(n: int) -> dict:
    """Simple compute task to run on cloud."""
    import math
    import time

    start = time.time()
    result = sum(math.sqrt(i) for i in range(n))
    elapsed = time.time() - start

    return {
        "n": n,
        "result": result,
        "elapsed": elapsed,
        "message": f"Computed sqrt sum of 0 to {n-1}"
    }


def prime_counter(limit: int) -> dict:
    """Count primes up to limit."""
    import time

    start = time.time()
    primes = []
    for n in range(2, limit + 1):
        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

    return {
        "limit": limit,
        "count": len(primes),
        "last_five": primes[-5:] if len(primes) >= 5 else primes,
        "elapsed": elapsed,
    }


def matrix_multiply(size: int) -> dict:
    """Matrix multiplication test."""
    # Create two random matrices
    import random
    import time
    random.seed(42)

    A = [[random.random() for _ in range(size)] for _ in range(size)]
    B = [[random.random() for _ in range(size)] for _ in range(size)]

    # Multiply them
    start = time.time()
    C = [[sum(A[i][k] * B[k][j] for k in range(size))
          for j in range(size)]
         for i in range(size)]
    elapsed = time.time() - start

    return {
        "size": size,
        "elapsed": elapsed,
        "sample_result": C[0][0] if size > 0 else None,
    }


# ============================================================================
# Test Runner
# ============================================================================

async def test_gcp_provider_connection():
    """Test 1: Verify GCP provider connection."""
    print_section("TEST 1: GCP Provider Connection")

    provider = GCPProvider()

    # Check credentials
    creds = provider.get_credentials()
    print(f"  Project ID: {creds.credentials.get('project_id')}")
    print(f"  Configured: {creds.is_configured}")
    print(f"  Free Tier Mode: {creds.credentials.get('free_tier_only')}")

    if not creds.is_configured:
        print(f"  ❌ FAILED: {creds.error_message}")
        return False

    # Verify connection
    success, error = await provider.verify_connection()
    if success:
        print("  ✅ Connection verified")
        return True
    else:
        print(f"  ❌ Connection failed: {error}")
        return False


async def test_function_serialization():
    """Test 2: Verify function serialization works."""
    print_section("TEST 2: Function Serialization")

    # Serialize test function
    try:
        serialized = cloudpickle.dumps(simple_compute)
        size = len(serialized)
        print(f"  Serialized simple_compute: {size} bytes")

        # Verify it can be deserialized
        restored = pickle.loads(serialized)
        result = restored(1000)
        print(f"  Executed restored function: result={result['result']:.2f}")

        # Test with all functions
        for func in [simple_compute, prime_counter, matrix_multiply]:
            serialized = cloudpickle.dumps(func)
            print(f"  {func.__name__}: {len(serialized)} bytes ✅")

        return True
    except Exception as e:
        print(f"  ❌ Serialization failed: {e}")
        return False


async def test_startup_script_generation():
    """Test 3: Generate startup script for cloud execution."""
    print_section("TEST 3: Startup Script Generation")

    # Create a startup script that runs our function
    func = simple_compute
    args = (100000,)
    kwargs = {}

    # Serialize function and arguments
    func_bytes = cloudpickle.dumps(func)
    args_bytes = cloudpickle.dumps(args)
    kwargs_bytes = cloudpickle.dumps(kwargs)

    import base64
    func_b64 = base64.b64encode(func_bytes).decode()
    args_b64 = base64.b64encode(args_bytes).decode()
    kwargs_b64 = base64.b64encode(kwargs_bytes).decode()

    startup_script = f'''#!/bin/bash
echo "=========================================="
echo "🚀 GateWay Job Starting"
echo "=========================================="
echo "Instance: $(hostname)"
echo "Date: $(date)"

# Install dependencies
apt-get update -qq
apt-get install -y -qq python3-pip > /dev/null 2>&1
pip3 install cloudpickle -q

# Run the job
python3 << 'PYTHON_SCRIPT'
import base64
import pickle
import cloudpickle
import json
import sys

print("Loading serialized function...")

# Decode and deserialize
func_bytes = base64.b64decode("{func_b64}")
args_bytes = base64.b64decode("{args_b64}")
kwargs_bytes = base64.b64decode("{kwargs_b64}")

func = pickle.loads(func_bytes)
args = pickle.loads(args_bytes)
kwargs = pickle.loads(kwargs_bytes)

print(f"Executing {{func.__name__}}{{args}}...")

try:
    result = func(*args, **kwargs)
    print("========================================")
    print("GATEWAY_RESULT_START")
    print(json.dumps(result, default=str))
    print("GATEWAY_RESULT_END")
    print("========================================")
    print("✅ Job completed successfully!")
except Exception as e:
    print(f"❌ Job failed: {{e}}")
    import traceback
    traceback.print_exc()
    sys.exit(1)
PYTHON_SCRIPT

echo ""
echo "=========================================="
echo "✅ GateWay Job Complete!"
echo "=========================================="

# Signal completion
touch /tmp/gateway_job_complete
'''

    print(f"  Script size: {len(startup_script)} bytes")
    print(f"  Function payload: {len(func_b64)} bytes (base64)")
    print("  ✅ Startup script generated")

    return startup_script


async def test_launch_and_execute(startup_script: str):
    """Test 4: Launch VM and execute job."""
    print_section("TEST 4: Launch VM & Execute Job")

    provider = GCPProvider()
    zone = os.environ.get("GCP_ZONE", "us-central1-a")

    instance = None
    try:
        # Launch instance
        print("  📦 Launching e2-micro instance...")
        instance = await provider.launch_instance(
            instance_type="e2-micro",
            region=zone,
            spot=False,  # Don't use preemptible for test
            startup_script=startup_script,
            tags={"test": "e2e", "purpose": "gateway-test", "managed": "true"},
        )

        print(f"  Instance ID: {instance.instance_id}")
        print(f"  Public IP: {instance.public_ip}")
        print(f"  Status: {instance.status}")

        # Wait for job to complete
        print("\n  ⏳ Waiting for job to complete...")
        max_wait = 180  # 3 minutes max
        waited = 0
        result_found = False
        job_result = None

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

            # Check serial output
            serial_output = provider.get_serial_output(instance.instance_id)

            if "GateWay Job Complete" in serial_output:
                result_found = True

                # Extract result
                if "GATEWAY_RESULT_START" in serial_output:
                    try:
                        start_idx = serial_output.index("GATEWAY_RESULT_START") + len("GATEWAY_RESULT_START")
                        end_idx = serial_output.index("GATEWAY_RESULT_END")
                        result_json = serial_output[start_idx:end_idx].strip()

                        # Clean up the JSON (remove log prefixes)
                        lines = result_json.split('\n')
                        for line in lines:
                            if line.strip().startswith('{'):
                                # Extract JSON from log line if needed
                                if 'command("/bin/bash"):' in line:
                                    line = line.split('command("/bin/bash"):')[-1].strip()
                                import json
                                job_result = json.loads(line)
                                break
                    except Exception as e:
                        print(f"  ⚠️  Could not parse result: {e}")
            else:
                print(f"     Still waiting... ({waited}s)")

        if result_found:
            print("\n  ✅ Job completed!")
            if job_result:
                print(f"  📊 Result: {job_result}")
            return True, job_result
        else:
            print(f"\n  ⚠️  Job didn't complete in {max_wait}s")
            return False, None

    finally:
        # Cleanup
        if instance:
            print(f"\n  🧹 Cleaning up instance {instance.instance_id}...")
            await provider.terminate_instance(instance.instance_id)
            print("  ✅ Instance terminated")


async def test_get_provider_manager():
    """Test 5: Test provider manager."""
    print_section("TEST 5: Provider Manager")

    # Test getting default provider
    provider = get_provider()
    print(f"  Default provider: {provider.name}")

    # Test getting GCP explicitly
    gcp = get_provider("gcp")
    print(f"  GCP provider: {gcp.name}")
    creds = gcp.get_credentials()
    print(f"  GCP configured: {creds.is_configured}")

    # Test listing regions
    regions = await gcp.list_regions()
    print(f"  Available regions: {regions}")

    # Test listing instance types
    types = await gcp.list_instance_types(region="us-central1-a")
    print(f"  Instance types: {[t['instance_type'] for t in types]}")

    return True


async def main():
    """Run all end-to-end tests."""

    # Safety check - don't run real tests by default
    if not REAL_TEST:
        print("=" * 60)
        print("⚠️  GateWay E2E 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 GATEWAY_E2E_TESTS=true")
        print("=" * 60)
        print("\n✅ Test skipped successfully (no resources created)")
        return 0

    print("\n" + "🚀" * 30)
    print("  GateWay End-to-End Test Suite")
    print("  ⚠️  REAL_TEST=True - This WILL create cloud resources!")
    print("🚀" * 30)

    results = {}

    # Test 1: Connection
    results["connection"] = await test_gcp_provider_connection()
    if not results["connection"]:
        print("\n❌ Cannot continue without GCP connection")
        return 1

    # Test 2: Serialization
    results["serialization"] = await test_function_serialization()

    # Test 3: Script generation
    startup_script = await test_startup_script_generation()
    results["script"] = startup_script is not None

    # Test 4: Launch and execute
    if startup_script:
        success, job_result = await test_launch_and_execute(startup_script)
        results["execution"] = success
    else:
        results["execution"] = False

    # Test 5: Provider manager
    results["provider_manager"] = await test_get_provider_manager()

    # Summary
    print_section("TEST SUMMARY")

    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" + "🎉" * 20)
        print("  ALL TESTS PASSED!")
        print("🎉" * 20)
        print("\n✅ GateWay is ready for cloud execution!")
        print("   Functions can be serialized, shipped to GCP, and executed.")
        print("   Results are retrieved from the cloud VM.")
    else:
        print("\n⚠️  Some tests failed")

    return 0 if all_passed else 1


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