#!/usr/bin/env python3
"""
Benchmark different LLM models for performance and quality comparison.
"""

import time
from typing import List, Dict
from pygeai import Client


def benchmark_model(
    client: Client,
    model: str,
    prompts: List[str],
    temperature: float = 0.7
) -> Dict:
    """Benchmark a single model."""
    
    results = {
        "model": model,
        "total_tokens": 0,
        "total_time": 0.0,
        "responses": []
    }
    
    for prompt in prompts:
        start = time.time()
        
        try:
            response = client.llm.complete(
                model=model,
                prompt=prompt,
                temperature=temperature
            )
            
            elapsed = time.time() - start
            
            results["responses"].append({
                "prompt": prompt[:50] + "...",
                "response_length": len(response.content),
                "tokens": response.usage.total_tokens,
                "time": elapsed
            })
            
            results["total_tokens"] += response.usage.total_tokens
            results["total_time"] += elapsed
            
        except Exception as e:
            print(f"Error with {model}: {e}")
            results["responses"].append({"error": str(e)})
    
    results["avg_time"] = results["total_time"] / len(prompts)
    results["tokens_per_second"] = results["total_tokens"] / results["total_time"]
    
    return results


def main():
    """Run benchmark comparison."""
    
    client = Client()
    
    # Test prompts
    prompts = [
        "Explain binary search in one paragraph",
        "Write a Python function to find prime numbers",
        "What are the SOLID principles?",
        "Compare REST vs GraphQL",
        "Explain async/await in JavaScript"
    ]
    
    # Models to compare
    models = [
        "gpt-4o",
        "gpt-4o-mini",
        "gpt-3.5-turbo",
        "claude-3-5-sonnet-20241022",
        "claude-3-5-haiku-20241022"
    ]
    
    print("🧪 Benchmarking Models\n")
    print("=" * 80)
    
    results = []
    
    for model in models:
        print(f"\nTesting {model}...")
        result = benchmark_model(client, model, prompts)
        results.append(result)
        
        print(f"  Avg Time: {result['avg_time']:.2f}s")
        print(f"  Tokens/sec: {result['tokens_per_second']:.1f}")
        print(f"  Total Tokens: {result['total_tokens']}")
    
    print("\n" + "=" * 80)
    print("\n📊 Summary\n")
    
    for result in sorted(results, key=lambda x: x["avg_time"]):
        print(f"{result['model']:30} | {result['avg_time']:.2f}s | {result['tokens_per_second']:.0f} tok/s")
    
    print("\n✅ Benchmark complete!")


if __name__ == "__main__":
    main()
