"""
Test Ollama API integration
"""
import requests
import json
import urllib3

# Disable SSL warnings
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

API_URL = "https://ollama_pjapi.theaken.com"

print("=" * 60)
print("Testing Ollama API Connection")
print("=" * 60)
print()

# Test 1: List models
print("Test 1: Listing available models...")
try:
    response = requests.get(f"{API_URL}/v1/models", timeout=10, verify=False)
    print(f"Status Code: {response.status_code}")

    if response.status_code == 200:
        data = response.json()
        models = data.get('data', [])
        print(f"Found {len(models)} models:")
        for model in models[:5]:
            print(f"  - {model.get('id', 'Unknown')}")
    else:
        print(f"Error: {response.text}")
except Exception as e:
    print(f"Error: {str(e)}")

print()

# Test 2: Chat completion
print("Test 2: Testing chat completion...")
try:
    chat_request = {
        "model": "qwen2.5:3b",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Say hello in Chinese."}
        ],
        "temperature": 0.7,
        "max_tokens": 50
    }

    response = requests.post(
        f"{API_URL}/v1/chat/completions",
        json=chat_request,
        headers={'Content-Type': 'application/json'},
        timeout=60,
        verify=False
    )

    print(f"Status Code: {response.status_code}")

    if response.status_code == 200:
        result = response.json()
        text = result['choices'][0]['message']['content']
        print(f"Response: {text}")
    else:
        print(f"Error: {response.text}")
except Exception as e:
    print(f"Error: {str(e)}")

print()
print("=" * 60)