import requests
import time
import json

# --- Configuration ---
OLLAMA_URL = "http://localhost:11434/api/generate"
LM_STUDIO_URL = "http://localhost:1234/v1/chat/completions"

OLLAMA_MODEL = "gemma4:e4b"
# LM Studio usually accepts the folder path or the filename as the model identifier.
LM_STUDIO_MODEL = "lmstudio-community/gemma-4-E4B-it-GGUF/gemma-4-E4B-it-Q8_0.gguf"

WARMUP_PROMPT = "Hi, respond with exactly one word: 'Ready'."
BENCHMARK_PROMPT = "Write a highly detailed, 5-paragraph essay on the history of the microprocessor, focusing on architectural changes from the 1970s to modern Apple Silicon."

ITERATIONS = 3 # How many times to run the benchmark per engine

def test_ollama(prompt, is_warmup=False):
    payload = {
        "model": OLLAMA_MODEL,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.0
        }
    }

    start_time = time.time()
    try:
        response = requests.post(OLLAMA_URL, json=payload)
        response.raise_for_status()
        end_time = time.time()

        data = response.json()
        duration = end_time - start_time

        # Ollama returns eval_count (completion tokens)
        tokens = data.get("eval_count", 0)
        tps = tokens / duration if duration > 0 else 0

        if not is_warmup:
            print(f"  -> Generated {tokens} tokens in {duration:.2f}s | Speed: {tps:.2f} tokens/sec")
        return tps
    except Exception as e:
        print(f"  -> Ollama error: {e}")
        return 0

def test_lm_studio(prompt, is_warmup=False):
    payload = {
        "model": LM_STUDIO_MODEL,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.0,
        "stream": False
    }

    start_time = time.time()
    try:
        response = requests.post(LM_STUDIO_URL, json=payload)
        response.raise_for_status()
        end_time = time.time()

        data = response.json()
        duration = end_time - start_time

        # LM Studio returns OpenAI-compatible usage stats
        tokens = data.get("usage", {}).get("completion_tokens", 0)
        tps = tokens / duration if duration > 0 else 0

        if not is_warmup:
            print(f"  -> Generated {tokens} tokens in {duration:.2f}s | Speed: {tps:.2f} tokens/sec")
        return tps
    except Exception as e:
        print(f"  -> LM Studio error: {e}")
        return 0

def main():
    print("========================================")
    print("🤖 LLM Engine Benchmark: Ollama vs LM Studio")
    print("========================================\n")

    # WARMUP PHASE
    print("--- Warming up models ---")
    print("Warming up Ollama...")
    test_ollama(WARMUP_PROMPT, is_warmup=True)
    print("Warming up LM Studio...")
    test_lm_studio(WARMUP_PROMPT, is_warmup=True)
    print("Warmup complete. Both models are in memory.\n")

    # BENCHMARK PHASE
    ollama_results = []
    lm_studio_results = []

    print(f"--- Starting Benchmark ({ITERATIONS} iterations) ---")
    for i in range(1, ITERATIONS + 1):
        print(f"\nIteration {i}/{ITERATIONS}:")

        print("Running Ollama...")
        tps_ollama = test_ollama(BENCHMARK_PROMPT)
        if tps_ollama > 0: ollama_results.append(tps_ollama)

        print("Running LM Studio...")
        tps_lm_studio = test_lm_studio(BENCHMARK_PROMPT)
        if tps_lm_studio > 0: lm_studio_results.append(tps_lm_studio)

    # RESULTS CALCULATION
    print("\n========================================")
    print("📊 Final Results (Average Tokens / Second)")
    print("========================================")

    if ollama_results:
        avg_ollama = sum(ollama_results) / len(ollama_results)
        print(f"Ollama Average:    {avg_ollama:.2f} t/s")
    else:
        print("Ollama Average:    Failed to run")

    if lm_studio_results:
        avg_lm_studio = sum(lm_studio_results) / len(lm_studio_results)
        print(f"LM Studio Average: {avg_lm_studio:.2f} t/s")
    else:
        print("LM Studio Average: Failed to run")

if __name__ == "__main__":
    main()