#!/usr/bin/env python3
"""
Configure Aiva (Genesis Prime) with comprehensive system prompt
"""

import requests
import json

LETTA_URL = "http://localhost:8283"
AGENT_ID = "agent-35039b24-fbb9-4139-b036-4a04e9c3d6ac"

def read_system_prompt():
    """Read the system prompt from file"""
    with open('/mnt/e/genesis-system/genesis_prime_system_prompt.txt', 'r') as f:
        return f.read()

def update_agent_system(system_prompt):
    """Update Aiva's system prompt"""

    print(f"🔧 Updating Aiva's system configuration...")
    print(f"   System prompt length: {len(system_prompt)} characters")

    try:
        # Update the agent's system prompt
        response = requests.patch(
            f"{LETTA_URL}/v1/agents/{AGENT_ID}",
            json={
                "system": system_prompt,
                "name": "aiva",  # Update name to Aiva
                "description": "Aiva - Genesis Prime Coordinator: Autonomous Master agent for intelligent task coordination and multi-agent orchestration"
            },
            headers={"Content-Type": "application/json"}
        )

        if response.status_code == 200:
            print(f"✅ Aiva's configuration updated successfully")
            return True
        else:
            print(f"❌ Failed to update Aiva: {response.status_code}")
            print(f"   Response: {response.text}")
            return False

    except Exception as e:
        print(f"❌ Error updating Aiva: {e}")
        return False

def test_aiva():
    """Test Aiva with a coordination query"""

    test_message = """Hello Aiva. This is Claude Code checking in for task assignment.

Current status:
- All infrastructure operational
- Ready to execute tasks
- Awaiting your coordination

What is my next priority task?"""

    print(f"\n🧪 Testing Aiva's coordination capabilities...")

    try:
        response = requests.post(
            f"{LETTA_URL}/v1/agents/{AGENT_ID}/messages",
            json={
                "messages": [{
                    "role": "user",
                    "content": test_message
                }]
            },
            headers={"Content-Type": "application/json"},
            timeout=120  # GLM can be slow
        )

        if response.status_code == 200:
            result = response.json()
            print(f"✅ Aiva responded successfully!")
            print(f"\n📨 Aiva's Response:")
            print("=" * 60)

            # Extract the assistant's response
            messages = result.get('messages', [])
            for msg in messages:
                role = msg.get('role', 'unknown')
                content = msg.get('content', '')
                if role == 'assistant' and content:
                    # Print full response
                    print(content)

            print("=" * 60)

            # Show token usage
            usage = result.get('usage', {})
            if usage:
                print(f"\n📊 Token Usage:")
                print(f"   Prompt: {usage.get('prompt_tokens', 'N/A')} tokens")
                print(f"   Completion: {usage.get('completion_tokens', 'N/A')} tokens")
                print(f"   Total: {usage.get('total_tokens', 'N/A')} tokens")

            return True
        else:
            print(f"❌ Aiva test failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return False

    except Exception as e:
        print(f"❌ Error testing Aiva: {e}")
        return False

def main():
    print("=" * 60)
    print("Aiva - Genesis Prime Coordinator Configuration")
    print("=" * 60)

    # Step 1: Read system prompt
    print("\n[1/3] Reading system prompt...")
    system_prompt = read_system_prompt()
    print(f"✅ System prompt loaded ({len(system_prompt.split())} words)")

    # Step 2: Update agent
    print("\n[2/3] Updating Aiva's configuration...")
    success = update_agent_system(system_prompt)

    if not success:
        print("\n❌ Configuration update failed")
        return 1

    # Step 3: Test
    print("\n[3/3] Testing Aiva's coordination capabilities...")
    test_success = test_aiva()

    # Summary
    print("\n" + "=" * 60)
    if test_success:
        print("✅ SUCCESS: Aiva is now fully configured as Genesis Prime!")
        print("\nAiva's Capabilities:")
        print("   • Autonomous 24/7 coordination")
        print("   • Intelligent task delegation")
        print("   • Persistent memory across sessions")
        print("   • Zero-cost operation (local LLM)")
        print("   • Master agent authority")
        print("\n🚀 Genesis Protocol coordination ready!")
        return 0
    else:
        print("⚠️  PARTIAL SUCCESS: Configuration updated but test inconclusive")
        print("   Aiva may need time to fully initialize")
        return 2

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