"""
Model Comparison Benchmark
Compare multiple models on the same prompts, measuring quality, latency, and cost.
"""
import os
import time
import json
from dataclasses import dataclass, asdict
from openai import OpenAI
import anthropic


@dataclass
class BenchmarkResult:
    model: str
    prompt_name: str
    response: str
    latency_ms: float
    input_tokens: int
    output_tokens: int
    estimated_cost: float


# ─── Configuration ───────────────────────────────────────────

OPENAI_MODELS = ["gpt-4o", "gpt-4o-mini"]
ANTHROPIC_MODELS = ["claude-sonnet-4-20250514", "claude-3-haiku-20240307"]

# Pricing per 1M tokens
PRICING = {
    "gpt-4o": {"input": 2.50, "output": 10.00},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "claude-sonnet-4-20250514": {"input": 3.00, "output": 15.00},
    "claude-3-haiku-20240307": {"input": 0.25, "output": 1.25},
}

# Test prompts — adjust these to match YOUR use case
TEST_PROMPTS = {
    "code_generation": {
        "system": "You are a Python expert. Write clean, well-commented code.",
        "user": "Write a function that finds the longest palindromic substring in a string. Include type hints and a docstring.",
    },
    "summarization": {
        "system": "You summarize text concisely.",
        "user": "Summarize in 2-3 sentences: Machine learning operations (MLOps) is a set of practices that aims to deploy and maintain machine learning models in production reliably and efficiently. The word is a compound of machine learning and the continuous development practice of DevOps. MLOps involves automating and monitoring all steps of ML system construction, including integration, testing, releasing, deployment, and infrastructure management.",
    },
    "reasoning": {
        "system": "You solve logic puzzles step by step.",
        "user": "If all roses are flowers, and some flowers fade quickly, can we conclude that some roses fade quickly? Explain your reasoning.",
    },
    "classification": {
        "system": "Classify the sentiment as positive, negative, or neutral. Reply with ONE word only.",
        "user": "The product works fine but the customer service was absolutely terrible and I waited 3 hours on hold.",
    },
}


# ─── Benchmark Functions ─────────────────────────────────────

def benchmark_openai(model: str, system: str, user: str) -> BenchmarkResult:
    """Benchmark a single OpenAI model call."""
    client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

    start = time.perf_counter()
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
        temperature=0,
        max_tokens=500,
    )
    latency = (time.perf_counter() - start) * 1000

    usage = response.usage
    cost = (
        (usage.prompt_tokens / 1_000_000) * PRICING[model]["input"]
        + (usage.completion_tokens / 1_000_000) * PRICING[model]["output"]
    )

    return BenchmarkResult(
        model=model,
        prompt_name="",
        response=response.choices[0].message.content,
        latency_ms=round(latency, 1),
        input_tokens=usage.prompt_tokens,
        output_tokens=usage.completion_tokens,
        estimated_cost=round(cost, 6),
    )


def benchmark_anthropic(model: str, system: str, user: str) -> BenchmarkResult:
    """Benchmark a single Anthropic model call."""
    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

    start = time.perf_counter()
    response = client.messages.create(
        model=model,
        max_tokens=500,
        system=system,
        messages=[{"role": "user", "content": user}],
        temperature=0,
    )
    latency = (time.perf_counter() - start) * 1000

    usage = response.usage
    cost = (
        (usage.input_tokens / 1_000_000) * PRICING[model]["input"]
        + (usage.output_tokens / 1_000_000) * PRICING[model]["output"]
    )

    return BenchmarkResult(
        model=model,
        prompt_name="",
        response=response.content[0].text,
        latency_ms=round(latency, 1),
        input_tokens=usage.input_tokens,
        output_tokens=usage.output_tokens,
        estimated_cost=round(cost, 6),
    )


def run_benchmark() -> list[BenchmarkResult]:
    """Run all benchmarks and return results."""
    results = []

    for prompt_name, prompt in TEST_PROMPTS.items():
        print(f"\n{'='*60}")
        print(f"  Testing: {prompt_name}")
        print(f"{'='*60}")

        for model in OPENAI_MODELS:
            print(f"  Running {model}...", end=" ", flush=True)
            result = benchmark_openai(model, prompt["system"], prompt["user"])
            result.prompt_name = prompt_name
            results.append(result)
            print(f"{result.latency_ms}ms, ${result.estimated_cost:.6f}")

        for model in ANTHROPIC_MODELS:
            print(f"  Running {model}...", end=" ", flush=True)
            result = benchmark_anthropic(model, prompt["system"], prompt["user"])
            result.prompt_name = prompt_name
            results.append(result)
            print(f"{result.latency_ms}ms, ${result.estimated_cost:.6f}")

    return results


def print_summary(results: list[BenchmarkResult]):
    """Print a summary table of results."""
    print(f"\n\n{'='*80}")
    print("  BENCHMARK SUMMARY")
    print(f"{'='*80}\n")

    # Group by prompt
    prompts = list(TEST_PROMPTS.keys())
    models = OPENAI_MODELS + ANTHROPIC_MODELS

    # Latency summary
    print(f"{'Model':<30} {'Avg Latency':<15} {'Avg Cost':<15} {'Avg Tokens Out'}")
    print("-" * 75)

    for model in models:
        model_results = [r for r in results if r.model == model]
        if not model_results:
            continue
        avg_latency = sum(r.latency_ms for r in model_results) / len(model_results)
        avg_cost = sum(r.estimated_cost for r in model_results) / len(model_results)
        avg_tokens = sum(r.output_tokens for r in model_results) / len(model_results)
        print(f"{model:<30} {avg_latency:>8.0f}ms     ${avg_cost:<12.6f} {avg_tokens:.0f}")

    # Save full results
    output_file = "benchmark_results.json"
    with open(output_file, "w") as f:
        json.dump([asdict(r) for r in results], f, indent=2)
    print(f"\nFull results saved to {output_file}")


if __name__ == "__main__":
    results = run_benchmark()
    print_summary(results)

# Install dependencies
pip install openai anthropic

# Set API keys
export OPENAI_API_KEY="sk-..."
export ANTHROPIC_API_KEY="sk-ant-..."

# Run benchmark
python model_benchmark.py