import subprocess
import time
import json
import os
from datetime import datetime

# Configuration
TEST_MODEL = "unsloth/gemma-4-E2B-it-GGUF"
PROMPTS_FILE = "../test/prompts.json"
RESULTS_FILE = "benchmark_results.md"
NUM_PROMPTS = 100 # Change this to test fewer prompts if needed

def load_prompts(file_path, limit=100):
    if not os.path.exists(file_path):
        print(f"Error: {file_path} not found.")
        return ["Write a short greeting."]
    with open(file_path, "r") as f:
        prompts = json.load(f)
    return prompts[:limit]

def benchmark_service(service_name, port, model_name, prompts):
    url = f"http://localhost:{port}/v1"
    
    print(f"\n--- Checking {service_name} on port {port} ---")
    try:
        res = subprocess.run(f"curl -s {url}/models", shell=True, capture_output=True, timeout=5)
        if res.returncode != 0:
            print(f"Skipping {service_name}: Service not found or not ready.")
            return None
    except Exception as e:
        print(f"Skipping {service_name}: Connection failed ({e})")
        return None

    print(f"Benchmarking {service_name} with {len(prompts)} prompts...")

    total_duration = 0
    total_tokens = 0
    results = []

    for i, prompt in enumerate(prompts):
        payload = {
            "model": model_name,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0
        }
        
        start_request = time.time()
        res = subprocess.run([
            "curl", "-s", f"{url}/chat/completions",
            "-H", "Content-Type: application/json",
            "-d", json.dumps(payload)
        ], capture_output=True, text=True)
        end_request = time.time()
        
        duration = end_request - start_request
        
        tokens = 0
        try:
            data = json.loads(res.stdout)
            tokens = data.get("usage", {}).get("completion_tokens", 0)
            if tokens == 0:
                 tokens = data.get("usage", {}).get("total_tokens", 0)
        except:
            print(f"Warning: Could not parse response for {service_name} on prompt {i+1}")
            continue

        total_duration += duration
        total_tokens += tokens
        tps = tokens / duration if duration > 0 else 0
        
        print(f"[{i+1}/{len(prompts)}] {duration:.2f}s, {tokens} tokens, {tps:.2f} tps")
        
    avg_duration = total_duration / len(prompts) if prompts else 0
    avg_tps = total_tokens / total_duration if total_duration > 0 else 0
    
    return {
        "service": service_name,
        "avg_duration": round(avg_duration, 2),
        "total_tokens": total_tokens,
        "avg_tps": round(avg_tps, 2),
        "num_prompts": len(prompts)
    }

def save_to_markdown(results):
    if not results or all(r is None for r in results):
        print("No results to save.")
        return

    header = f"\n## Benchmark Results (100 Prompts) - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
    header += "| Service | Avg Latency (s) | Total Tokens | Avg Tokens/Sec | Prompts |\n"
    header += "| :--- | :--- | :--- | :--- | :--- |\n"
    
    rows = ""
    for r in results:
        if r:
            rows += f"| {r['service']} | {r['avg_duration']}s | {r['total_tokens']} | {r['avg_tps']} | {r['num_prompts']} |\n"
    
    with open(RESULTS_FILE, "a") as f:
        f.write(header + rows + "\n")
    print(f"\nResults saved to {RESULTS_FILE}")

if __name__ == "__main__":
    prompts = load_prompts(PROMPTS_FILE, NUM_PROMPTS)
    all_results = []
    
    # Test vLLM (Port 8021)
    vllm_res = benchmark_service("vllm-gemma", 8021, "google/gemma-4-E2B", prompts)
    if vllm_res: all_results.append(vllm_res)
    
    # Test llama.cpp (Port 8088)
    llamacpp_res = benchmark_service("llamacpp-gemma", 8088, "unsloth/gemma-4-E2B-it-GGUF", prompts)
    if llamacpp_res: all_results.append(llamacpp_res)
    
    save_to_markdown(all_results)