#!/usr/bin/env python3
"""
grade_run.py — run a prompt against LM Studio's local server, save the raw
output, capture metrics from the API, and draft a benchmark_history.json entry.

Solves two problems at once:
  1. No clipboard → outputs aren't mangled in transit (the file is written
     directly from the API response).
  2. Metrics (tok/sec, total tokens, TTFT) come from the API response — no
     manual entry.

USAGE
  python3 tools/grade_run.py \\
      --lmstudio http://10.0.0.31:1234 \\
      --model "qwen3.6-35b-a3b-6bit-mlx" \\
      --prompt tts \\
      --name qwen3.6-35b-a3b-6bit-mlx-tts

WHAT IT DOES
  1. Reads prompts/<prompt>_<word>.txt (lfu_cache, tts_pipeline, webhook_bridge,
     rust_service, data_service, automation_glue, mcp_server).
  2. POSTs to {lmstudio}/v1/chat/completions with the model id.
  3. Saves the raw completion to outputs/<name>.<ext>  (ext inferred from prompt).
  4. Prints a DRAFT JSON entry (you/me fill the 5-pillar scores + audit by hand).
  5. Optionally --append writes a placeholder entry to data/benchmark_history.json
     with total_score:null, verdict:"pending", so the audit is the only manual bit.

NOTE: This only COLLECTS facts (runnability + metrics + file). The actual
5-pillar audit is still done by the grader (me/you) — that part is subjective
and shouldn't be faked.
"""
import argparse, json, os, sys, time, urllib.request, urllib.error

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # repo root
PROMPTS = {
    "lfu":        ("prompts/lfu_cache_prompt.txt", "py"),
    "tts":        ("prompts/tts_pipeline.txt",     "py"),
    "webhook":    ("prompts/webhook_bridge.txt",   "py"),
    "data":       ("prompts/data_service.txt",     "py"),
    "automation": ("prompts/automation_glue.txt",  "py"),
    "mcp":        ("prompts/mcp_server.txt",       "ts"),   # or py
    "rust":       ("prompts/rust_service.txt",     "rs"),
}

def api(base, path, payload=None, timeout=600):
    url = base.rstrip("/") + path
    data = json.dumps(payload).encode() if payload else None
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST" if payload else "GET")
    t0 = time.monotonic()
    with urllib.request.urlopen(req, timeout=timeout) as r:
        body = json.loads(r.read())
    return body, time.monotonic() - t0

def list_models(base):
    try:
        body, _ = api(base, "/v1/models")
        return [m.get("id") for m in body.get("data", [])]
    except Exception as e:
        return []

def main():
    ap = argparse.ArgumentParser(description="Run a prompt via LM Studio, save output + metrics.")
    ap.add_argument("--lmstudio", default="http://10.0.0.31:1234", help="LM Studio base URL")
    ap.add_argument("--model", help="Model id as LM Studio reports it (use --list to see)")
    ap.add_argument("--prompt", choices=list(PROMPTS), help="Which prompt to run")
    ap.add_argument("--name", help="Output filename stem, e.g. qwen3.6-35b-a3b-6bit-mlx-tts")
    ap.add_argument("--list", action="store_true", help="Just list loaded models and exit")
    ap.add_argument("--append", action="store_true", help="Append a draft (pending) entry to data/benchmark_history.json")
    ap.add_argument("--max-tokens", type=int, default=8192, help="Max completion tokens")
    args = ap.parse_args()

    # --list works standalone; otherwise model/prompt/name are required
    if not args.list and not (args.model and args.prompt and args.name):
        ap.error("--model, --prompt, and --name are required (unless --list)")

    if args.list:
        ms = list_models(args.lmstudio)
        print("Models loaded in LM Studio:")
        for m in ms: print("  -", m)
        if not ms: print("  (none / server not reachable)")
        return

    prompt_file, ext = PROMPTS[args.prompt]
    ppath = os.path.join(HERE, prompt_file)
    if not os.path.exists(ppath):
        sys.exit(f"prompt file not found: {ppath}")
    prompt_text = open(ppath).read()

    print(f"→ POST {args.lmstudio}/v1/chat/completions  model={args.model}  prompt={args.prompt}")
    payload = {
        "model": args.model,
        "messages": [{"role": "user", "content": prompt_text}],
        "temperature": 0.2,
        "max_tokens": args.max_tokens,
        "stream": False,
    }
    try:
        resp, wall = api(args.lmstudio, "/v1/chat/completions", payload)
    except urllib.error.URLError as e:
        sys.exit(f"LM Studio not reachable at {args.lmstudio} — is the server started and on 0.0.0.0? ({e})")

    content = resp["choices"][0]["message"]["content"]
    usage = resp.get("usage", {})
    # LM Studio / OpenAI-style usage: prompt_tokens, completion_tokens, total_tokens
    comp_tokens = usage.get("completion_tokens") or usage.get("completion_tokens_details", {}).get("reasoning_tokens", 0)
    ttft = resp.get("stats", {}).get("time_to_first_token") or resp.get("timings", {}).get("prompt_n")
    tok_sec = resp.get("stats", {}).get("tokens_per_second") or resp.get("timings", {}).get("predicted_per_second")
    # fallback: derive tok/sec from wall time if the server didn't report it
    if tok_sec is None and comp_tokens and wall:
        tok_sec = round(comp_tokens / wall, 2)

    # save the raw output
    out_path = os.path.join(HERE, "outputs", f"{args.name}.{ext}")
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    # strip common markdown fences if present
    text = content.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        if lines[0].startswith("```"): lines = lines[1:]
        if lines and lines[-1].strip() == "```": lines = lines[:-1]
        text = "\n".join(lines)
    open(out_path, "w").write(text)
    print(f"✓ saved {len(text)} bytes -> {os.path.relpath(out_path, HERE)}")

    print("\n=== METRICS (from API) ===")
    print(f"  tok/sec      : {tok_sec}")
    print(f"  total tokens : {usage.get('total_tokens')}")
    print(f"  TTFT         : {ttft}")

    print("\n=== DRAFT JSON ENTRY (fill breakdown + audit by hand) ===")
    entry = {
        "id": args.name,
        "prompt_id": args.prompt,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "model_name": args.model,
        "quant": "TODO",
        "format": "mlx",
        "tok_sec": tok_sec,
        "total_tokens": usage.get("total_tokens"),
        "ttft_sec": ttft,
        "filename": f"outputs/{args.name}.{ext}",
        "tests_pass": None,  # grader runs the file
        "total_score": None,
        "breakdown": {},
        "verdict": "pending",
        "best_for": "TODO",
        "critical_bugs": [],
        "patch_code": "",
    }
    print(json.dumps(entry, indent=2))

    if args.append:
        hp = os.path.join(HERE, "data", "benchmark_history.json")
        d = json.load(open(hp))
        # replace if same id+prompt exists, else append
        d["models"] = [m for m in d["models"] if not (m["id"] == entry["id"] and m["prompt_id"] == entry["prompt_id"])]
        d["models"].append(entry)
        json.dump(d, open(hp, "w"), indent=2)
        print(f"\n✓ appended pending entry to {os.path.relpath(hp, HERE)} (run the audit, then fill breakdown/total_score/verdict)")

if __name__ == "__main__":
    main()