import os
import asyncio
from openai import AsyncOpenAI, BadRequestError
from dataclasses import dataclass
from typing import Callable

client = AsyncOpenAI(
    api_key=os.environ["RADIUM_API_KEY"],
    base_url="https://api.radium.cloud/v1",
)

MODELS = ["tycho-1.0", "clarke-1.0", "hal-1.0"]


@dataclass
class TrialResult:
    model: str
    content: str
    latency_ms: float
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


async def run_trial(prompt: str, model: str) -> TrialResult:
    import time
    start = time.perf_counter()
    try:
        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.7,
            max_tokens=512,
        )
    except BadRequestError as e:
        return TrialResult(
            model=model,
            content=f"ERROR: {e}",
            latency_ms=0,
            prompt_tokens=0,
            completion_tokens=0,
            total_tokens=0,
        )

    latency = (time.perf_counter() - start) * 1000
    usage = response.usage
    return TrialResult(
        model=model,
        content=response.choices[0].message.content,
        latency_ms=latency,
        prompt_tokens=usage.prompt_tokens,
        completion_tokens=usage.completion_tokens,
        total_tokens=usage.total_tokens,
    )


def score_length(result: TrialResult) -> float:
    """Prefer concise answers (higher score = worse)."""
    return len(result.content)


def score_formatting(result: TrialResult) -> float:
    """Prefer answers with bullet points or numbered lists."""
    return 1.0 if any(marker in result.content for marker in ["- ", "1. ", "* "]) else 0.0


def print_comparison(results: list[TrialResult]):
    print(f"\n{'Model':<15} {'Latency':>10} {'Prompt':>8} {'Completion':>12} {'Total':>8}")
    print("-" * 60)
    for r in results:
        print(
            f"{r.model:<15} "
            f"{r.latency_ms:>9.0f}ms "
            f"{r.prompt_tokens:>8} "
            f"{r.completion_tokens:>12} "
            f"{r.total_tokens:>8}"
        )

    print("\n--- Outputs ---")
    for r in results:
        print(f"\n[{r.model}]:\n{r.content[:300]}...")


async def main():
    prompt = (
        "Explain how HTTP/3 works and why it's faster than HTTP/2. "
        "Use bullet points."
    )

    print(f"Prompt: {prompt}\n")
    results = await asyncio.gather(*[run_trial(prompt, m) for m in MODELS])

    # Custom scoring
    for scorer in [score_length, score_formatting]:
        results.sort(key=scorer, reverse=(scorer.__name__ == "score_formatting"))
        winner = results[0]
        print(f"\n🏆 Winner by {scorer.__name__}: {winner.model}")

    print_comparison(results)


if __name__ == "__main__":
    asyncio.run(main())