#!/usr/bin/env python3
"""
Agent Registration Script for WHOOSH
Registers cluster agents with the WHOOSH orchestration system
"""

import json
import requests
import yaml
import sys
import time
from pathlib import Path

# Configuration
WHOOSH_API_URL = "https://whoosh.home.deepblack.cloud/api"
CONFIG_FILE = "/home/tony/AI/projects/whoosh/config/whoosh.yaml"

def load_config():
    """Load the whoosh.yaml configuration file"""
    try:
        with open(CONFIG_FILE, 'r') as f:
            return yaml.safe_load(f)
    except Exception as e:
        print(f"❌ Error loading config: {e}")
        sys.exit(1)

def test_whoosh_connection():
    """Test connection to WHOOSH API"""
    try:
        response = requests.get(f"{WHOOSH_API_URL}/health", timeout=5)
        if response.status_code == 200:
            print("✅ Connected to WHOOSH API")
            return True
        else:
            print(f"❌ WHOOSH API returned status {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Failed to connect to WHOOSH API: {e}")
        return False

def test_agent_connectivity(endpoint):
    """Test if an agent endpoint is responsive"""
    try:
        response = requests.get(f"{endpoint}/api/tags", timeout=5)
        return response.status_code == 200
    except:
        return False

def register_agent(agent_id, agent_config):
    """Register a single agent with WHOOSH"""
    
    # Check if agent is responsive
    if not test_agent_connectivity(agent_config['endpoint']):
        print(f"⚠️  {agent_id.upper()} is not responsive at {agent_config['endpoint']}")
        return False
    
    # Prepare agent registration data
    agent_data = {
        "id": agent_id,
        "endpoint": agent_config['endpoint'],
        "model": agent_config['model'],
        "specialty": agent_config['specialization'],
        "capabilities": agent_config['capabilities'],
        "hardware": agent_config['hardware'],
        "performance_targets": agent_config['performance_targets'],
        "status": "available",
        "current_tasks": 0,
        "max_concurrent": 3  # Default concurrent task limit
    }
    
    try:
        # Register the agent
        response = requests.post(
            f"{WHOOSH_API_URL}/api/agents",
            json=agent_data,
            headers={"Content-Type": "application/json"},
            timeout=10
        )
        
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Registered {agent_id.upper()} - Agent ID: {result.get('agent_id', 'Unknown')}")
            return True
        else:
            print(f"❌ Failed to register {agent_id.upper()}: {response.status_code} - {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ Error registering {agent_id.upper()}: {e}")
        return False

def main():
    """Main registration process"""
    print("🐝 WHOOSH Agent Registration Script")
    print("=" * 50)
    
    # Test WHOOSH connection
    if not test_whoosh_connection():
        print("❌ Cannot connect to WHOOSH API. Make sure WHOOSH is running.")
        sys.exit(1)
    
    # Load configuration
    config = load_config()
    agents = config.get('whoosh', {}).get('agents', {})
    
    if not agents:
        print("❌ No agents found in configuration")
        sys.exit(1)
    
    print(f"📋 Found {len(agents)} agents to register:")
    for agent_id in agents.keys():
        print(f"  • {agent_id.upper()}")
    
    print("\n🔄 Starting registration process...")
    
    # Register each agent
    successful_registrations = 0
    failed_registrations = 0
    
    for agent_id, agent_config in agents.items():
        print(f"\n📡 Registering {agent_id.upper()}...")
        
        if register_agent(agent_id, agent_config):
            successful_registrations += 1
        else:
            failed_registrations += 1
        
        time.sleep(1)  # Brief pause between registrations
    
    # Summary
    print("\n" + "=" * 50)
    print(f"📊 Registration Summary:")
    print(f"  ✅ Successful: {successful_registrations}")
    print(f"  ❌ Failed: {failed_registrations}")
    print(f"  📈 Total: {successful_registrations + failed_registrations}")
    
    if successful_registrations > 0:
        print(f"\n🎉 Successfully registered {successful_registrations} agents!")
        print("🔗 Check agent status: curl https://whoosh.home.deepblack.cloud/api/agents")
    else:
        print("\n💔 No agents were successfully registered.")
        sys.exit(1)

if __name__ == "__main__":
    main()