#!/usr/bin/env python3
"""
Configure Letta Genesis agent to use GLM-4.6V-Flash via LM Studio
"""

import requests
import json
import sys
from time import sleep

import subprocess

LETTA_URL = "http://localhost:8283"
# Get Windows host IP for WSL
WINDOWS_HOST = subprocess.check_output(
    "ip route | grep default | awk '{print $3}'",
    shell=True
).decode().strip()
LMSTUDIO_URL = f"http://{WINDOWS_HOST}:1234"
AGENT_ID = "agent-35039b24-fbb9-4139-b036-4a04e9c3d6ac"

def check_lmstudio():
    """Check if LM Studio is running and get available models"""
    try:
        response = requests.get(f"{LMSTUDIO_URL}/v1/models", timeout=5)
        if response.status_code == 200:
            models = response.json()
            print(f"✅ LM Studio is running at {LMSTUDIO_URL}")
            print(f"📋 Available models:")
            for model in models.get('data', []):
                print(f"   - {model.get('id', 'unknown')}")
            return models.get('data', [])
        else:
            print(f"⚠️  LM Studio responded with status {response.status_code}")
            return None
    except requests.exceptions.ConnectionError:
        print(f"❌ LM Studio not responding at {LMSTUDIO_URL}")
        print("   Please ensure LM Studio is running with local server enabled")
        return None
    except Exception as e:
        print(f"❌ Error checking LM Studio: {e}")
        return None

def get_agent_config():
    """Get current Genesis agent configuration"""
    try:
        response = requests.get(f"{LETTA_URL}/v1/agents/{AGENT_ID}")
        if response.status_code == 200:
            return response.json()
        else:
            print(f"❌ Failed to get agent config: {response.status_code}")
            return None
    except Exception as e:
        print(f"❌ Error getting agent config: {e}")
        return None

def update_agent_llm(model_name="glm-4.6v-flash"):
    """Update Genesis agent to use LM Studio LLM"""

    # Prepare the new LLM configuration
    new_llm_config = {
        "model": model_name,
        "model_endpoint_type": "openai",  # LM Studio uses OpenAI-compatible API
        "model_endpoint": f"{LMSTUDIO_URL}/v1",
        "context_window": 32768,  # GLM-4.6V-Flash context window
        "put_inner_thoughts_in_kwargs": False,
        "temperature": 0.7,
        "max_tokens": None,
        "parallel_tool_calls": False
    }

    # Prepare new embedding config (use LM Studio if available, else keep lightweight)
    new_embedding_config = {
        "embedding_endpoint_type": "openai",
        "embedding_endpoint": f"{LMSTUDIO_URL}/v1",
        "embedding_model": "nomic-embed-text",  # Common in LM Studio
        "embedding_dim": 768,  # Standard for nomic-embed
        "embedding_chunk_size": 300,
        "batch_size": 32
    }

    print(f"\n🔧 Updating Genesis agent configuration...")
    print(f"   Model: {model_name}")
    print(f"   Endpoint: {LMSTUDIO_URL}/v1")
    print(f"   Context: {new_llm_config['context_window']} tokens")

    try:
        # Update the agent
        response = requests.patch(
            f"{LETTA_URL}/v1/agents/{AGENT_ID}",
            json={
                "llm_config": new_llm_config,
                "embedding_config": new_embedding_config
            },
            headers={"Content-Type": "application/json"}
        )

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

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

def test_agent():
    """Test Genesis agent with local LLM"""
    test_message = "Hello Genesis Core. This is a test message from Claude Code. Please confirm you can receive this and respond intelligently using the local GLM-4.6V-Flash model."

    print(f"\n🧪 Testing agent communication...")

    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=60  # GLM might be slower than GPT-4
        )

        if response.status_code == 200:
            result = response.json()
            print(f"✅ Agent responded successfully!")
            print(f"\n📨 Response:")

            # Extract and display the response
            messages = result.get('messages', [])
            for msg in messages:
                role = msg.get('role', 'unknown')
                content = msg.get('content', '')
                if role == 'assistant' and content:
                    print(f"   {content[:500]}...")  # First 500 chars

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

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

def main():
    print("=" * 60)
    print("Letta Genesis Agent Configuration for LM Studio")
    print("=" * 60)

    # Step 1: Check LM Studio
    print("\n[1/4] Checking LM Studio availability...")
    models = check_lmstudio()

    if models is None:
        print("\n⚠️  LM Studio is not available yet.")
        print("   Please:")
        print("   1. Ensure GLM-4.6V-Flash download is complete")
        print("   2. Start the local server in LM Studio (port 1234)")
        print("   3. Load the GLM-4.6V-Flash model")
        print("   4. Re-run this script")
        return 1

    # Detect model name
    model_name = "glm-4.6v-flash"  # Default
    if models:
        # Try to find GLM model
        for model in models:
            model_id = model.get('id', '')
            if 'glm' in model_id.lower():
                model_name = model_id
                print(f"✅ Detected GLM model: {model_name}")
                break

    # Step 2: Get current config
    print("\n[2/4] Getting current agent configuration...")
    current_config = get_agent_config()
    if current_config:
        print(f"✅ Current model: {current_config.get('llm_config', {}).get('model', 'unknown')}")

    # Step 3: Update configuration
    print("\n[3/4] Updating agent to use LM Studio...")
    success = update_agent_llm(model_name)

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

    # Step 4: Test
    print("\n[4/4] Testing agent with local LLM...")
    sleep(2)  # Give system a moment to apply changes
    test_success = test_agent()

    # Summary
    print("\n" + "=" * 60)
    if test_success:
        print("✅ SUCCESS: Genesis agent is now using GLM-4.6V-Flash!")
        print(f"   Model: {model_name}")
        print(f"   Endpoint: {LMSTUDIO_URL}/v1")
        print("\n🚀 Letta can now coordinate intelligently using the local LLM")
        print("   This means FREE 24/7 operation with no API costs!")
        return 0
    else:
        print("⚠️  PARTIAL SUCCESS: Configuration updated but test failed")
        print("   The agent may need time to initialize")
        print("   Try sending a message manually to verify")
        return 2

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