#!/usr/bin/env python3
"""
Test script for Portainer Edge MCP Server
Tests edge environments, stacks, groups, and jobs functionality
"""

import asyncio
import aiohttp
import os
import json

PORTAINER_URL = os.getenv("PORTAINER_URL", "").rstrip("/")
PORTAINER_API_KEY = os.getenv("PORTAINER_API_KEY", "")

async def make_request(method: str, endpoint: str, json_data=None, params=None):
    """Make a request to Portainer API"""
    url = f"{PORTAINER_URL}{endpoint}"
    headers = {"X-API-Key": PORTAINER_API_KEY}
    
    async with aiohttp.ClientSession() as session:
        async with session.request(
            method,
            url,
            json=json_data,
            params=params,
            headers=headers
        ) as response:
            text = await response.text()
            if response.status >= 400:
                print(f"❌ Error {response.status}: {text}")
                return None
            return json.loads(text) if text else {}

async def test_edge_environments():
    """Test edge environment operations"""
    print("\n🌐 Testing Edge Environments...")
    
    # List all endpoints/environments
    endpoints = await make_request("GET", "/api/endpoints")
    if endpoints:
        edge_envs = [e for e in endpoints if e.get("Type") == 4]  # Type 4 is Edge
        print(f"✅ Found {len(edge_envs)} edge environments")
        
        if edge_envs:
            # Get details of first edge environment
            env = edge_envs[0]
            print(f"  • {env['Name']} (ID: {env['Id']})")
            print(f"    Status: {env.get('Status', 'Unknown')}")
            print(f"    Edge ID: {env.get('EdgeID', 'N/A')}")
            
            # Get edge status
            status = await make_request("GET", f"/api/endpoints/{env['Id']}/edge/status")
            if status:
                print(f"    Check-in: {status.get('CheckinTime', 'Never')}")
    else:
        print("❌ Failed to list environments")

async def test_edge_groups():
    """Test edge group operations"""
    print("\n👥 Testing Edge Groups...")
    
    # List edge groups
    groups = await make_request("GET", "/api/edge_groups")
    if groups:
        print(f"✅ Found {len(groups)} edge groups")
        for group in groups[:3]:  # Show first 3
            print(f"  • {group['Name']} (ID: {group['Id']})")
            print(f"    Dynamic: {'Yes' if group.get('Dynamic') else 'No'}")
            print(f"    Endpoints: {len(group.get('Endpoints', []))}")
    else:
        print("❌ Failed to list edge groups")
    
    # Try to create a test edge group
    print("\n📝 Creating test edge group...")
    test_group_data = {
        "Name": "test-edge-group",
        "Dynamic": False,
        "TagIds": []
    }
    
    new_group = await make_request("POST", "/api/edge_groups", json_data=test_group_data)
    if new_group:
        print(f"✅ Created edge group: {new_group['Name']} (ID: {new_group['Id']})")
        
        # Clean up - delete the test group
        await make_request("DELETE", f"/api/edge_groups/{new_group['Id']}")
        print("🗑️ Cleaned up test edge group")
    else:
        print("❌ Failed to create edge group")

async def test_edge_stacks():
    """Test edge stack operations"""
    print("\n📚 Testing Edge Stacks...")
    
    # List edge stacks
    stacks = await make_request("GET", "/api/edge_stacks")
    if stacks:
        print(f"✅ Found {len(stacks)} edge stacks")
        for stack in stacks[:3]:  # Show first 3
            print(f"  • {stack['Name']} (ID: {stack['Id']})")
            print(f"    Type: {stack.get('StackType', 'Unknown')}")
            print(f"    Groups: {len(stack.get('EdgeGroups', []))}")
            
            # Check if it has GitOps
            if stack.get("GitConfig") and stack.get("AutoUpdate"):
                print(f"    GitOps: Enabled ({stack['AutoUpdate'].get('Interval', 'N/A')})")
    else:
        print("❌ Failed to list edge stacks")

async def test_edge_jobs():
    """Test edge job operations"""
    print("\n💼 Testing Edge Jobs...")
    
    # List edge jobs
    jobs = await make_request("GET", "/api/edge_jobs")
    if jobs:
        print(f"✅ Found {len(jobs)} edge jobs")
        for job in jobs[:3]:  # Show first 3
            print(f"  • {job['Name']} (ID: {job['Id']})")
            print(f"    Recurring: {'Yes' if job.get('Recurring') else 'No'}")
            if job.get('CronExpression'):
                print(f"    Schedule: {job['CronExpression']}")
            print(f"    Target Groups: {len(job.get('EdgeGroups', []))}")
    else:
        print("❌ Failed to list edge jobs")

async def test_edge_settings():
    """Test edge settings"""
    print("\n⚙️ Testing Edge Settings...")
    
    # Get settings
    settings = await make_request("GET", "/api/settings")
    if settings and settings.get("Edge"):
        edge_settings = settings["Edge"]
        print("✅ Edge Settings:")
        print(f"  • Check-in Interval: {edge_settings.get('CheckinInterval', 'N/A')} seconds")
        print(f"  • Command Interval: {edge_settings.get('CommandInterval', 'N/A')} seconds")
        print(f"  • Ping Interval: {edge_settings.get('PingInterval', 'N/A')} seconds")
        print(f"  • Tunnel Server: {edge_settings.get('TunnelServerAddress', 'Not configured')}")
    else:
        print("❌ Failed to get edge settings")

async def main():
    """Run all edge tests"""
    print("🚀 Portainer Edge API Tests")
    print(f"URL: {PORTAINER_URL}")
    print(f"API Key: {'***' + PORTAINER_API_KEY[-4:] if PORTAINER_API_KEY else 'Not set'}")
    
    if not PORTAINER_URL or not PORTAINER_API_KEY:
        print("\n❌ Please set PORTAINER_URL and PORTAINER_API_KEY environment variables")
        return
    
    # Run tests
    await test_edge_environments()
    await test_edge_groups()
    await test_edge_stacks()
    await test_edge_jobs()
    await test_edge_settings()
    
    print("\n✅ Edge API tests completed!")

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