#!/usr/bin/env python3
"""
Test script to verify environment variable passing to agent pods
"""
import requests
import json
import sys

API_URL = "http://localhost:8000"

def test_create_mysql_agent_with_env():
    """Test creating a MySQL agent with environment variables"""
    
    payload = {
        "name": "test-mysql-agent",
        "template": "mysql_agent",
        "env": {
            "MYSQL_HOST": "test-mysql-server.mysql.database.azure.com",
            "MYSQL_PORT": "3306",
            "MYSQL_USER": "testuser",
            "MYSQL_PASSWORD": "testpass",
            "MYSQL_DATABASE": "testdb",
            "OPENAI_API_KEY": "sk-test-key"
        },
        "config": {
            "cpu_request": "100m",
            "memory_request": "128Mi"
        }
    }
    
    print("Creating MySQL agent with environment variables...")
    print(f"Payload: {json.dumps(payload, indent=2)}")
    
    try:
        response = requests.post(f"{API_URL}/agents", json=payload)
        print(f"\nStatus Code: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        
        if response.status_code == 200:
            print("\n✅ Agent created successfully!")
            return True
        else:
            print("\n❌ Agent creation failed!")
            return False
            
    except Exception as e:
        print(f"\n❌ Error: {str(e)}")
        return False

def test_get_agent_status():
    """Test getting agent status"""
    
    agent_name = "test-mysql-agent"
    
    print(f"\nGetting status for agent: {agent_name}")
    
    try:
        response = requests.get(f"{API_URL}/agents/{agent_name}/status")
        print(f"Status Code: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        
        if response.status_code == 200:
            print("\n✅ Status retrieved successfully!")
            return True
        else:
            print("\n❌ Failed to get status!")
            return False
            
    except Exception as e:
        print(f"\n❌ Error: {str(e)}")
        return False

def test_get_template_info():
    """Test getting template information"""
    
    template = "mysql_agent"
    
    print(f"\nGetting template info for: {template}")
    
    try:
        response = requests.get(f"{API_URL}/templates/{template}")
        print(f"Status Code: {response.status_code}")
        print(f"Response: {json.dumps(response.json(), indent=2)}")
        
        if response.status_code == 200:
            print("\n✅ Template info retrieved successfully!")
            return True
        else:
            print("\n❌ Failed to get template info!")
            return False
            
    except Exception as e:
        print(f"\n❌ Error: {str(e)}")
        return False

def cleanup_test_agent():
    """Cleanup test agent"""
    
    agent_name = "test-mysql-agent"
    
    print(f"\nCleaning up test agent: {agent_name}")
    
    try:
        response = requests.delete(f"{API_URL}/agents/{agent_name}")
        print(f"Status Code: {response.status_code}")
        
        if response.status_code == 200:
            print("✅ Test agent cleaned up!")
        else:
            print("⚠️  Cleanup may have failed (agent might not exist)")
            
    except Exception as e:
        print(f"⚠️  Cleanup error: {str(e)}")

if __name__ == "__main__":
    print("=" * 60)
    print("Testing Agent Manager API - Environment Variables")
    print("=" * 60)
    
    # Test 1: Get template info
    test_get_template_info()
    
    print("\n" + "=" * 60)
    
    # Test 2: Create agent with env vars
    test_create_mysql_agent_with_env()
    
    print("\n" + "=" * 60)
    
    # Wait a bit for pod to be created
    import time
    print("\nWaiting 5 seconds for pod creation...")
    time.sleep(5)
    
    # Test 3: Get agent status
    test_get_agent_status()
    
    print("\n" + "=" * 60)
    
    # Cleanup
    cleanup_test_agent()
    
    print("\n" + "=" * 60)
    print("Tests completed!")
    print("=" * 60)