#!/usr/bin/env python3
"""Benchmark LLMHosts Core vs Ollama on equivalent models.

Measures: time-to-first-token (TTFT), tok/s, total latency, memory usage.
"""

import json
import sys
import time


def benchmark_core(model_path: str, prompt_tokens: list[int], n_tokens: int = 50) -> dict:
    from llmhosts.inference.core import LLMHostsInferenceEngine

    engine = LLMHostsInferenceEngine()

    t0 = time.perf_counter()
    engine.load_model(model_path)
    load_time = time.perf_counter() - t0

    t0 = time.perf_counter()
    tokens = engine.generate_tokens(prompt_tokens, max_tokens=n_tokens, temperature=0.0)
    gen_time = time.perf_counter() - t0

    return {
        "engine": "llmhosts",
        "load_time_s": round(load_time, 3),
        "gen_time_s": round(gen_time, 3),
        "tokens_generated": len(tokens),
        "tok_per_s": round(len(tokens) / gen_time, 1) if gen_time > 0 else 0,
    }


if __name__ == "__main__":
    model = sys.argv[1] if len(sys.argv) > 1 else None
    if not model:
        print("Usage: benchmark_core.py <model.gguf>")
        sys.exit(1)
    result = benchmark_core(model, [1, 5, 10, 20, 30], n_tokens=50)
    print(json.dumps(result, indent=2))
