#!/usr/bin/env python3
"""
Simple test script to verify Ollama is working
"""
import requests
import json

def test_ollama_connection():
    """Test if Ollama is running and accessible"""
    try:
        response = requests.get("http://localhost:11434/api/tags", timeout=10)
        if response.status_code == 200:
            models = response.json().get("models", [])
            print(f"✅ Ollama is running!")
            print(f"📋 Available models: {[m['name'] for m in models]}")
            return True
        else:
            print(f"❌ Ollama health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ Cannot connect to Ollama: {e}")
        return False

def test_ollama_generation():
    """Test if Ollama can generate text"""
    try:
        payload = {
            "model": "llama3.2:latest",
            "prompt": "Hello! Please respond with 'Ollama is working correctly!'",
            "stream": False
        }
        
        response = requests.post(
            "http://localhost:11434/api/generate",
            json=payload,
            timeout=30
        )
        
        if response.status_code == 200:
            result = response.json()
            generated_text = result.get('response', '').strip()
            print(f"✅ Ollama generation test successful!")
            print(f"🤖 Response: {generated_text}")
            return True
        else:
            print(f"❌ Ollama generation failed: {response.status_code} - {response.text}")
            return False
            
    except Exception as e:
        print(f"❌ Ollama generation test failed: {e}")
        return False

if __name__ == "__main__":
    print("🧪 Testing Ollama Setup...")
    print("=" * 50)
    
    # Test connection
    if test_ollama_connection():
        print()
        # Test generation
        test_ollama_generation()
    
    print("=" * 50)
    print("✅ Ollama test completed!")