#!/usr/bin/env python3
"""StudioLLM Comprehensive Benchmark Suite"""
import json
import time
import subprocess
import sys
import os

BASE = "https://inference.studio.dripper.live/api/ollama/v1"
AUTH = "ollama"

MODELS = [
    "ornith:35b", "ornith:9b",
    "llama3.2:3b", "moondream:latest", "llama3.2-vision:11b", "llava:7b",
    "minicpm-v:latest", "deepseek-r1:14b", "qwen2.5:32b",
    "deepseek-coder-v2:latest", "deepseek-r1:8b", "qwen2.5-coder:7b",
    "phi3:mini", "mistral:7b", "llama3.1:8b", "qwen2.5:14b", "qwen2.5:3b"
]

RESULTS = {}
WARMED = []

def curl_post(model, prompt="say hello in 3 words", max_tokens=20, timeout=120, temp=0.0, stream=False):
    """Send a request to StudioLLM and return timing + response."""
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": max_tokens,
        "temperature": temp,
        "stream": stream
    }
    cmd = [
        "curl", "-s", "-w", "\n%{http_code}",
        "--connect-timeout", "10",
        "--max-time", str(timeout),
        "-H", f"Authorization: Bearer {AUTH}",
        "-H", "Content-Type: application/json",
        "-d", json.dumps(payload),
        f"{BASE}/chat/completions"
    ]
    start = time.time()
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout+10)
    elapsed = time.time() - start
    
    if result.returncode != 0:
        return {"error": f"curl failed: {result.stderr.strip()}", "elapsed": elapsed, "http_code": 0}
    
    parts = result.stdout.strip().rsplit("\n", 1)
    http_code = int(parts[-1]) if parts[-1].isdigit() else 0
    body = parts[0] if len(parts) > 1 else ""
    
    try:
        data = json.loads(body) if body else {}
    except:
        data = {"raw": body[:500] if body else ""}
    
    return {"data": data, "elapsed": elapsed, "http_code": http_code, "raw_length": len(body)}

def get_first_token_time(model, timeout=120):
    """Stream a response and measure time to first token."""
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": "Write a 2-sentence story about a cat."}],
        "max_tokens": 50,
        "temperature": 0.0,
        "stream": True
    }
    cmd = [
        "curl", "-s", "-N",
        "--connect-timeout", "10",
        "--max-time", str(timeout),
        "-H", f"Authorization: Bearer {AUTH}",
        "-H", "Content-Type: application/json",
        "-d", json.dumps(payload),
        f"{BASE}/chat/completions"
    ]
    start = time.time()
    first_token = None
    total_content = ""
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout+10)
        total_time = time.time() - start
        if result.returncode == 0:
            for line in result.stdout.strip().split("\n"):
                if line.startswith("data: ") and line != "data: [DONE]":
                    try:
                        chunk = json.loads(line[6:])
                        choices = chunk.get("choices", [])
                        if choices and choices[0].get("delta", {}).get("content"):
                            if first_token is None:
                                first_token = time.time() - start
                            total_content += choices[0]["delta"]["content"]
                    except:
                        pass
    except Exception as e:
        total_time = time.time() - start
    
    return {
        "first_token_s": round(first_token, 2) if first_token else None,
        "total_time_s": round(total_time, 2) if first_token else round(total_time, 2),
        "response_length": len(total_content)
    }

def warmup(model):
    """Warm up a model - load into GPU memory."""
    print(f"  Warming {model}...", end=" ", flush=True)
    r = curl_post(model, "reply exactly: ready", max_tokens=10, timeout=120)
    if r.get("http_code") == 200 and "error" not in r.get("data", {}):
        print(f"OK ({r['elapsed']:.1f}s)")
        return True, r
    else:
        err = r.get("data", {}).get("error", r.get("error", "unknown"))
        print(f"FAIL - {str(err)[:80]}")
        return False, r

def benchmark_simple(model):
    """Simple response latency benchmark."""
    print(f"  Simple prompt...", end=" ", flush=True)
    r = curl_post(model, "What is 2+2? Answer with just the number.", max_tokens=10, timeout=120)
    elapsed = r["elapsed"]
    content = ""
    if r.get("http_code") == 200:
        content = r["data"].get("choices", [{}])[0].get("message", {}).get("content", "")
    print(f"{elapsed:.1f}s | {content.strip()[:50]!r}")
    return {"elapsed": round(elapsed, 2), "content": content.strip()[:100], "http": r["http_code"]}

def benchmark_reasoning(model):
    """Reasoning quality benchmark."""
    prompt = "A snail climbs a 30-foot wall. Each day it climbs 3 feet, but each night it slips back 2 feet. How many days does it take to reach the top? Show your work."
    print(f"  Reasoning test...", end=" ", flush=True)
    r = curl_post(model, prompt, max_tokens=500, timeout=120)
    elapsed = r["elapsed"]
    content = ""
    if r.get("http_code") == 200:
        content = r["data"].get("choices", [{}])[0].get("message", {}).get("content", "")
        # Check for reasoning_content
        reasoning = r["data"].get("choices", [{}])[0].get("message", {}).get("reasoning_content", "")
    print(f"{elapsed:.1f}s | {len(content)} chars")
    has_reasoning = "reasoning" in content.lower() or "step" in content.lower() or "day" in content.lower()
    return {
        "elapsed": round(elapsed, 2),
        "content_len": len(content),
        "has_reasoning_steps": has_reasoning,
        "content_preview": content[:200] if content else ""
    }

def benchmark_multiturn(model):
    """Multi-turn conversation test."""
    print(f"  Multi-turn...")
    turns = [
        "What is the capital of France?",
        "What is the population of that city?",
        "Name 3 famous landmarks there."
    ]
    history = []
    turn_results = []
    for i, q in enumerate(turns):
        msgs = history + [{"role": "user", "content": q}]
        payload = {
            "model": model,
            "messages": msgs,
            "max_tokens": 100,
            "temperature": 0.0
        }
        cmd = [
            "curl", "-s", "-w", "\n%{http_code}",
            "--connect-timeout", "10",
            "--max-time", "120",
            "-H", f"Authorization: Bearer {AUTH}",
            "-H", "Content-Type: application/json",
            "-d", json.dumps(payload),
            f"{BASE}/chat/completions"
        ]
        start = time.time()
        res = subprocess.run(cmd, capture_output=True, text=True, timeout=130)
        elapsed = time.time() - start
        parts = res.stdout.strip().rsplit("\n", 1)
        http = int(parts[-1]) if parts[-1].isdigit() else 0
        body = parts[0] if len(parts) > 1 else ""
        try:
            data = json.loads(body) if body else {}
        except:
            data = {}
        content = data.get("choices", [{}])[0].get("message", {}).get("content", "") if http == 200 else ""
        history.append({"role": "user", "content": q})
        if content:
            history.append({"role": "assistant", "content": content})
        turn_results.append({"turn": i+1, "elapsed": round(elapsed, 2), "http": http, "len": len(content)})
        print(f"    Turn {i+1}: {elapsed:.1f}s, {len(content)} chars -> {content[:60]!r}")
    return {"turns": turn_results, "total_turns": len(turn_results)}

def benchmark_tool_format(model):
    """Test JSON/Tool-calling format adherence."""
    prompt = """You are a weather assistant. Return a JSON object with this exact structure:
{"location": "San Francisco", "temperature": 18, "unit": "celsius", "condition": "foggy"}
Return ONLY valid JSON, no other text."""
    print(f"  Tool/JSON format...", end=" ", flush=True)
    r = curl_post(model, prompt, max_tokens=200, timeout=120, temp=0.1)
    content = ""
    elapsed = r["elapsed"]
    if r.get("http_code") == 200:
        content = r["data"].get("choices", [{}])[0].get("message", {}).get("content", "")
    # Check if valid JSON
    is_json = False
    try:
        parsed = json.loads(content)
        is_json = isinstance(parsed, dict) and "location" in parsed
    except:
        # Try to find JSON in response
        import re
        m = re.search(r'\{[^}]+\}', content)
        if m:
            try:
                parsed = json.loads(m.group())
                is_json = isinstance(parsed, dict) and "location" in parsed
            except:
                pass
    print(f"{elapsed:.1f}s | JSON: {'YES' if is_json else 'NO'}")
    return {"elapsed": round(elapsed, 2), "valid_json": is_json, "content": content[:150]}

def benchmark_long_context(model):
    """Progressive context test."""
    print(f"  Long context...")
    results = []
    for tokens in [500, 2000, 5000]:
        filler = "The quick brown fox jumps over the lazy dog. " * (tokens // 10)
        prompt = f"{filler}\n\nWhat was the last sentence about? Answer in 3 words."
        r = curl_post(model, prompt, max_tokens=20, timeout=120)
        elapsed = r["elapsed"]
        content = ""
        if r.get("http_code") == 200:
            content = r["data"].get("choices", [{}])[0].get("message", {}).get("content", "")
        status = "OK" if r["http_code"] == 200 else f"FAIL({r['http_code']})"
        print(f"    ~{tokens} token context: {elapsed:.1f}s, {status}, resp: {content[:40]!r}")
        results.append({"context_tokens": tokens, "elapsed": round(elapsed, 2), "http": r["http_code"]})
    return results

def benchmark_concurrent(model):
    """Test concurrent requests."""
    print(f"  Concurrent (3x)...")
    import threading
    results = []
    lock = threading.Lock()
    def send(idx):
        r = curl_post(model, f"What is {idx} + {idx*2}? Answer with just the number.", max_tokens=10, timeout=180)
        with lock:
            results.append({"idx": idx, "elapsed": round(r["elapsed"], 2), "http": r["http_code"]})
    threads = [threading.Thread(target=send, args=(i,)) for i in range(3)]
    start = time.time()
    for t in threads: t.start()
    for t in threads: t.join()
    total = time.time() - start
    statuses = [r["http"] for r in results]
    avg_time = sum(r["elapsed"] for r in results) / len(results) if results else 0
    all_ok = all(s == 200 for s in statuses)
    print(f"    Total: {total:.1f}s, Avg: {avg_time:.1f}s, All OK: {all_ok}")
    return {"concurrent_results": results, "total_wall_time": round(total, 2), "all_ok": all_ok, "avg_latency": round(avg_time, 2)}

def main():
    # Phase 1: Warmup all models
    print("=" * 70)
    print("PHASE 1: WARMUP ALL MODELS")
    print("=" * 70)
    for model in MODELS:
        ok, r = warmup(model)
        if ok:
            WARMED.append(model)
    
    print(f"\nWarmed up: {len(WARMED)}/{len(MODELS)} models")
    
    # Phase 2: Full benchmarks on responsive models
    print("\n" + "=" * 70)
    print("PHASE 2: FULL BENCHMARKS")
    print("=" * 70)
    
    for model in WARMED:
        print(f"\n--- {model} ---")
        model_results = {}
        
        # Skip vision models for text-only tests
        is_vision = any(v in model for v in ["vision", "moondream", "llava", "minicpm"])
        
        try:
            model_results["simple"] = benchmark_simple(model)
        except Exception as e:
            model_results["simple"] = {"error": str(e)}
        
        try:
            model_results["first_token"] = get_first_token_time(model)
        except Exception as e:
            model_results["first_token"] = {"error": str(e)}
        
        try:
            model_results["reasoning"] = benchmark_reasoning(model)
        except Exception as e:
            model_results["reasoning"] = {"error": str(e)}
        
        try:
            model_results["multiturn"] = benchmark_multiturn(model)
        except Exception as e:
            model_results["multiturn"] = {"error": str(e)}
        
        try:
            model_results["tool_format"] = benchmark_tool_format(model)
        except Exception as e:
            model_results["tool_format"] = {"error": str(e)}
        
        try:
            model_results["long_context"] = benchmark_long_context(model)
        except Exception as e:
            model_results["long_context"] = {"error": str(e)}
        
        try:
            model_results["concurrent"] = benchmark_concurrent(model)
        except Exception as e:
            model_results["concurrent"] = {"error": str(e)}
        
        RESULTS[model] = model_results
    
    # Phase 3: Generate report
    print("\n" + "=" * 70)
    print("GENERATING REPORT")
    print("=" * 70)
    
    report_path = sys.argv[1] if len(sys.argv) > 1 else "/home/amu/.openclaw-mustafagateway/agents/mustafa/workspace/.openclaw/tmp/studiollm-benchmark.md"
    
    with open(report_path, "w") as f:
        f.write(f"# StudioLLM Benchmark Report\n\n")
        f.write(f"**Date:** {time.strftime('%Y-%m-%d %H:%M UTC')}\n\n")
        
        f.write("## Models Available\n\n")
        f.write(f"**Total on server:** {len(MODELS)}\n")
        f.write(f"**Registered in Gateway:** ornith:9b, ornith:35b\n\n")
        f.write("| Model | Warmup | Reason |\n")
        f.write("|-------|--------|--------|\n")
        for m in MODELS:
            warmed = "✅" if m in WARMED else "❌"
            f.write(f"| {m} | {warmed} | |\n")
        
        f.write("\n## Benchmark Results\n\n")
        
        # Rank by speed
        ranked = []
        for model, res in RESULTS.items():
            simple = res.get("simple", {})
            rsn = res.get("reasoning", {})
            tool = res.get("tool_format", {})
            ft = res.get("first_token", {})
            speed = simple.get("elapsed", 999)
            first_token = ft.get("first_token_s", speed)
            quality = "Good"
            if rsn.get("has_reasoning_steps"):
                quality = "Excellent" if res.get("tool_format", {}).get("valid_json") else "Good"
            else:
                quality = "Poor reasoning"
            conc = res.get("concurrent", {})
            ranked.append({
                "model": model,
                "speed": speed,
                "first_token": first_token,
                "quality": quality,
                "reasoning": rsn.get("has_reasoning_steps", False),
                "json_ok": tool.get("valid_json", False),
                "concurrent_ok": conc.get("all_ok", False) if isinstance(conc, dict) else False,
                "full": res
            })
        
        ranked.sort(key=lambda x: x["speed"])
        
        f.write("### Speed Ranking\n\n")
        f.write("| Rank | Model | Simple (s) | First Token (s) | Reasoning | JSON | Concurrent |\n")
        f.write("|------|-------|-----------|-----------------|-----------|------|------------|\n")
        for i, r in enumerate(ranked, 1):
            rsn = "✅" if r["reasoning"] else "❌"
            js = "✅" if r["json_ok"] else "❌"
            conc = "✅" if r["concurrent_ok"] else "❌"
            f.write(f"| {i} | {r['model']} | {r['speed']} | {r['first_token']} | {rsn} | {js} | {conc} |\n")
        
        f.write("\n### Detailed Model Results\n\n")
        for model, res in RESULTS.items():
            f.write(f"#### {model}\n\n")
            simple = res.get("simple", {})
            f.write(f"- **Simple:** {simple.get('elapsed', 'ERR')}s | {simple.get('content', '')}\n")
            ft = res.get("first_token", {})
            f.write(f"- **First Token:** {ft.get('first_token_s', 'N/A')}s | Total: {ft.get('total_time_s', 'N/A')}s\n")
            rsn = res.get("reasoning", {})
            f.write(f"- **Reasoning:** {rsn.get('elapsed', 'ERR')}s | {rsn.get('content_len', 0)} chars | Steps: {'✅' if rsn.get('has_reasoning_steps') else '❌'}\n")
            mt = res.get("multiturn", {})
            if isinstance(mt, dict) and "turns" in mt:
                for t in mt["turns"]:
                    f.write(f"  - Turn {t['turn']}: {t['elapsed']}s, {t['len']} chars\n")
            tool = res.get("tool_format", {})
            f.write(f"- **JSON Format:** {'✅' if tool.get('valid_json') else '❌'} ({tool.get('elapsed', 'ERR')}s)\n")
            lc = res.get("long_context", [])
            if isinstance(lc, list):
                for ctx in lc:
                    f.write(f"  - {ctx.get('context_tokens', '?')} tokens: {ctx.get('elapsed', 'ERR')}s HTTP {ctx.get('http')}\n")
            conc = res.get("concurrent", {})
            if isinstance(conc, dict) and "concurrent_results" in conc:
                f.write(f"- **Concurrent:** Wall: {conc.get('total_wall_time', '?')}s | Avg: {conc.get('avg_latency', '?')}s | All OK: {'✅' if conc.get('all_ok') else '❌'}\n")
            f.write("\n")
        
        f.write("## Recommendations\n\n")
        f.write("### Fast Chat (low latency, simple tasks)\n\n")
        for r in ranked:
            if r["speed"] < 10:
                f.write(f"- **{r['model']}** — {r['speed']}s simple, first token {r['first_token']}s\n")
        
        f.write("\n### Heavy Reasoning (math, logic, analysis)\n\n")
        for r in ranked:
            if r["reasoning"]:
                f.write(f"- **{r['model']}** — {r['speed']}s simple, {r['quality']}\n")
        
        f.write("\n### Agent Assignment Recommendations\n\n")
        f.write("| Agent | Current Model | Recommended | Reason |\n")
        f.write("|-------|--------------|-------------|--------|\n")
        f.write("| mustafa | deepseek-v4-flash (external) | See below | Default agent, needs reliability |\n")
        f.write("| nexus | deepseek-v4-flash (external) | See below | Heavy research agent |\n")
        f.write("| dave | ornith:9b → ornith:35b → phi3:mini | Keep as-is or adjust | Already configured with StudioLLM |\n")
        f.write("| susie | deepseek-v4-flash (external) | See below | Light tasks |\n")
        f.write("| archie | deepseek-v4-pro (external) | See below | Heavy agent, needs best reasoning |\n")
        
        f.write("\n### Gateway Config Notes\n\n")
        f.write(f"- Only **ornith:9b** and **ornith:35b** registered in Gateway models\n")
        f.write(f"- `injectNumCtxForOpenAICompat: false` — needs verification with models that expect num_ctx\n")
        f.write(f"- `timeoutSeconds: 300` in provider config should let big models load\n")
        f.write(f"- `thinkingFormat: qwen-chat-template` — reasoning models should return `reasoning_content`\n")
        f.write("- All 17 server models should be registered in Gateway to allow model fallback chains\n")
    
    print(f"\nReport written to: {report_path}")
    print(f"\nRaw data also saved for analysis")
    
    # Also save raw results
    raw_path = report_path.replace(".md", "-raw.json")
    with open(raw_path, "w") as f:
        # Remove large content fields for JSON
        clean = {}
        for model, res in RESULTS.items():
            clean[model] = res
        json.dump(clean, f, indent=2, default=str)
    
    return report_path

if __name__ == "__main__":
    main()