#!/usr/bin/env python3
"""Small, explicitly non-ranking sample of each model's documented thinking defaults."""

from __future__ import annotations

import argparse
import json
import statistics
import time
from pathlib import Path

from harness_lib import Server, dump, post

ROOT = Path("/workspace/kingy-qwen38")
PROMPTS = [
    {
        "id": "decision",
        "prompt": "A 24GB GPU must serve a 32K coding workspace with a 4K output reserve. Explain a cautious process for selecting a quantization and KV cache. End with a one-sentence recommendation.",
    },
    {
        "id": "coding",
        "prompt": "Review this Python bug and propose a corrected function plus three tests: def merge_intervals(items): return sorted(items). The intended behavior merges overlapping or touching integer intervals without mutating input.",
    },
    {
        "id": "knowledge_work",
        "prompt": "Write a concise internal memo distinguishing a model that loads, a model that runs one short request, and a configuration that is comfortable for a three-turn long-context session. Do not invent measurements.",
    },
]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--model-path", type=Path, required=True)
    parser.add_argument("--reasoning-effort", default="default")
    parser.add_argument("--top-k", type=int, required=True)
    parser.add_argument("--root", type=Path, default=ROOT)
    parser.add_argument("--server-bin", type=Path, default=ROOT / "llama.cpp/build/bin/llama-server")
    args = parser.parse_args()
    out = args.root / "results" / args.model_id
    server = Server(args.model_id, args.model_path, args.server_bin, out, ctx=8192, reasoning="on", effort=args.reasoning_effort, label="native-default").start()
    rows = []
    try:
        for case in PROMPTS:
            payload = {
                "model": args.model_id,
                "messages": [{"role": "user", "content": case["prompt"]}],
                "seed": 20260816,
                "temperature": 1.0,
                "top_p": 0.95,
                "top_k": args.top_k,
                "min_p": 0.0,
                "presence_penalty": 0.0,
                "repeat_penalty": 1.0,
                "max_tokens": 1024,
                "stream": False,
            }
            began = time.perf_counter()
            response = post("/v1/chat/completions", payload)
            rows.append({"id": case["id"], "elapsed_seconds": time.perf_counter() - began, "request": payload, "response": response})
    finally:
        server_summary = server.stop()
    completion_tokens = [int((row["response"].get("usage") or {}).get("completion_tokens") or 0) for row in rows]
    elapsed = [row["elapsed_seconds"] for row in rows]
    dump(out / "native-default-summary.json", {
        "model_id": args.model_id,
        "documented_sampling": {"temperature": 1.0, "top_p": 0.95, "top_k": args.top_k, "min_p": 0.0, "presence_penalty": 0.0, "repetition_penalty": 1.0},
        "reasoning_effort": args.reasoning_effort,
        "ranking_use": "Non-ranking user-experience sample only",
        "rows": rows,
        "median_elapsed_seconds": statistics.median(elapsed),
        "completion_tokens_total": sum(completion_tokens),
        "server": server_summary,
    })
    print(json.dumps({"model": args.model_id, "median_elapsed_seconds": statistics.median(elapsed), "completion_tokens_total": sum(completion_tokens)}), flush=True)


if __name__ == "__main__":
    main()
