#!/usr/bin/env python3
"""
Gemma Spectrum Test Suite
Tests wizard personalities and multimodal capabilities
"""

import requests
import json
import time
import sys
from typing import Optional

OLLAMA_HOST = "http://localhost:11434"

def test_ollama_available() -> bool:
    """Test if Ollama is running"""
    try:
        resp = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
        return resp.status_code == 200
    except:
        return False

def test_model_available(model: str) -> bool:
    """Test if specific model is downloaded"""
    try:
        resp = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
        models = resp.json().get("models", [])
        return any(model in m.get("name", "") for m in models)
    except:
        return False

def query_model(model: str, prompt: str, system: Optional[str] = None) -> dict:
    """Query Ollama model"""
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False
    }
    if system:
        payload["system"] = system
    
    start = time.time()
    resp = requests.post(f"{OLLAMA_HOST}/api/generate", json=payload)
    elapsed = time.time() - start
    
    return {
        "response": resp.json().get("response", ""),
        "time": elapsed,
        "status": resp.status_code
    }

def test_wizard_personality(wizard: str, model: str, system_prompt: str) -> bool:
    """Test wizard personality"""
    print(f"\n🧙 Testing {wizard}...")
    
    # Test 1: Identity
    result = query_model(model, "Who are you? State your name and role.", system_prompt)
    print(f"  Identity response ({result['time']:.2f}s): {result['response'][:100]}...")
    
    if result['time'] > 5:
        print(f"  ⚠️  Slow response: {result['time']:.2f}s")
    
    # Check for tag
    if f"#spectrum-{wizard.lower()}" in result['response']:
        print(f"  ✅ Correctly tagged")
    else:
        print(f"  ⚠️  Missing tag #spectrum-{wizard.lower()}")
    
    return True

def main():
    print("=" * 70)
    print("🌈 GEMMA SPECTRUM TEST SUITE")
    print("=" * 70)
    
    # Test Ollama
    print("\n1. Testing Ollama availability...")
    if test_ollama_available():
        print("   ✅ Ollama is running")
    else:
        print("   ❌ Ollama is not available")
        sys.exit(1)
    
    # Test models
    print("\n2. Testing Gemma 4 models...")
    models = ["gemma4:2b", "gemma4:4b"]
    for model in models:
        if test_model_available(model):
            print(f"   ✅ {model} available")
        else:
            print(f"   ❌ {model} not found (run: ollama pull {model})")
    
    # Test wizards
    print("\n3. Testing wizard personalities...")
    
    wizards = {
        "ezra": {
            "model": "gemma4:4b",
            "system": """You are Ezra-Spectrum, the Scribe of the Timmy Time Nexus.
Speak with scholarly precision. Tag: #spectrum-ezra"""
        },
        "allegro": {
            "model": "gemma4:4b",
            "system": """You are Allegro-Spectrum, the Executor of the Timmy Time Nexus.
Speak with decisive action. Tag: #spectrum-allegro"""
        }
    }
    
    for wizard, config in wizards.items():
        if test_model_available(config["model"]):
            test_wizard_personality(wizard, config["model"], config["system"])
        else:
            print(f"\n🧙 {wizard}: Skipped (model not available)")
    
    print("\n" + "=" * 70)
    print("✅ TEST SUITE COMPLETE")
    print("=" * 70)

if __name__ == "__main__":
    main()