#!/usr/bin/env python3
"""
Test connection to 4080 LLM Server.
"""

import requests
import json
from config import OLLAMA_BASE_URL, API_TAGS, API_CHAT, MODEL_NAME

def test_server_connection():
    """Test if Ollama server is reachable."""
    print(f"Testing connection to {OLLAMA_BASE_URL}...")
    
    try:
        # Test tags endpoint
        response = requests.get(API_TAGS, timeout=5)
        if response.status_code == 200:
            data = response.json()
            print(f"✅ Server is reachable!")
            print(f"Available models: {len(data.get('models', []))}")
            for model in data.get('models', []):
                print(f"  - {model.get('name', 'unknown')}")
            return True
        else:
            print(f"❌ Server returned status {response.status_code}")
            return False
    except requests.exceptions.ConnectionError:
        print(f"❌ Cannot connect to {OLLAMA_BASE_URL}")
        print("   Make sure the server is running and accessible")
        return False
    except Exception as e:
        print(f"❌ Error: {e}")
        return False

def test_chat():
    """Test chat endpoint with a simple prompt."""
    print(f"\nTesting chat endpoint with model: {MODEL_NAME}...")
    
    payload = {
        "model": MODEL_NAME,
        "messages": [
            {"role": "user", "content": "Say 'Hello from 4080!' in one sentence."}
        ],
        "stream": False
    }
    
    try:
        response = requests.post(API_CHAT, json=payload, timeout=60)
        if response.status_code == 200:
            data = response.json()
            message = data.get('message', {})
            content = message.get('content', '')
            print(f"✅ Chat test successful!")
            print(f"Response: {content}")
            return True
        else:
            print(f"❌ Chat test failed: {response.status_code}")
            print(f"Response: {response.text}")
            return False
    except Exception as e:
        print(f"❌ Chat test error: {e}")
        return False

if __name__ == "__main__":
    print("=" * 60)
    print("4080 LLM Server Connection Test")
    print("=" * 60)
    
    if test_server_connection():
        test_chat()
    else:
        print("\n⚠️  Server connection failed. Check:")
        print("   1. Server is running on the GPU VM")
        print("   2. Network connectivity to 10.0.30.63:11434")
        print("   3. Firewall allows connections")