#!/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.

STATS COME FROM THE NATIVE API, NOT OPENAI-COMPAT.
  LM Studio's OpenAI-compatible endpoint (/v1/chat/completions) returns an EMPTY
  stats object — use the native /api/v1/chat endpoint instead, which returns a
  full stats block including time_to_first_token_seconds. Request shape differs
  from OpenAI: pass {"model":..., "input": "<prompt text>", ...} (input is a
  string or array of input items, NOT messages).

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, method=None):
    url = base.rstrip("/") + path
    if method is None:
        method = "POST" if payload is not None else "GET"
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method=method)
    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 []

# ---- memory-safe load/unload (NEVER have 2 big models resident at once) ----
def resident_instances(base):
    """Return list of (model_key, instance_id) for models actually in RAM.
    Uses /api/v1/models -> loaded_instances (the ONLY reliable resident list;
    /v1/models lists available models, not resident ones). Read-only, loads nothing."""
    try:
        body, _ = api(base, "/api/v1/models", None)
    except urllib.error.HTTPError as e:
        # endpoint may need GET not POST; fall back
        body, _ = api(base, "/api/v1/models")
    ms = body.get("data") or body.get("models") or []
    out = []
    for m in ms:
        for inst in m.get("loaded_instances", []):
            out.append((m.get("key") or m.get("id"), inst.get("id")))
    return out

def unload(base, instance_id):
    """POST /api/v1/models/unload {instance_id}. Returns True on success."""
    try:
        body, _ = api(base, "/api/v1/models/unload", {"instance_id": instance_id}, timeout=60)
        return body.get("instance_id") == instance_id or isinstance(body, dict)
    except Exception as e:
        print(f"  unload FAILED for {instance_id}: {e}", file=sys.stderr)
        return False

def unload_all(base, except_id=None):
    """Unload every resident instance except optionally one. Verify after."""
    resid = resident_instances(base)
    did = []
    for key, iid in resid:
        if except_id and (iid == except_id or key == except_id):
            continue
        if unload(base, iid):
            did.append(iid)
        else:
            print(f"  WARNING: could not unload {iid}", file=sys.stderr)
    # verify
    resid2 = resident_instances(base)
    return did, resid2

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")
    ap.add_argument("--unload-all", action="store_true", help="Unload ALL resident models, verify, and exit (no run)")
    ap.add_argument("--resident", action="store_true", help="Show what's actually resident in RAM and exit (no run)")
    args = ap.parse_args()

    # standalone info/safety commands first
    if args.unload_all:
        print("=== unloading all resident models ===")
        before = resident_instances(args.lmstudio)
        print("  before:", [i for _, i in before] or "(none)")
        did, after = unload_all(args.lmstudio)
        print("  unloaded:", did or "(nothing to unload)")
        print("  after (resident):", [i for _, i in after] or "(none — clean)")
        return
    if args.resident:
        r = resident_instances(args.lmstudio)
        print("Resident in RAM right now:")
        for key, iid in r: print(f"  {key}  (instance {iid})")
        if not r: print("  (none)")
        return
    if args.list:
        ms = list_models(args.lmstudio)
        print("Models available in LM Studio:")
        for m in ms: print("  -", m)
        if not ms: print("  (none / server not reachable)")
        return

    # --list/resident/unload-all work standalone; otherwise model/prompt/name required
    if not (args.model and args.prompt and args.name):
        ap.error("--model, --prompt, and --name are required (or use --list/--resident/--unload-all)")

    # ---- MEMORY GUARD: refuse to stack a second big model ----
    resid = resident_instances(args.lmstudio)
    resident_ids = [iid for _, iid in resid]
    if resid and args.model not in resident_ids and not any(args.model in (iid or "") or args.model == k for k, iid in resid):
        print(f"!! ABORT: a DIFFERENT model is resident ({resid}) but you asked to run '{args.model}'.", file=sys.stderr)
        print(f"!! Unload it first: python3 tools/grade_run.py --unload-all   (or load '{args.model}' and unload others)", file=sys.stderr)
        sys.exit(1)

    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()

    # NATIVE endpoint /api/v1/chat (not OpenAI-compat). The OpenAI endpoint
    # /v1/chat/completions returns an EMPTY stats object; the native endpoint
    # returns full stats including time_to_first_token_seconds. Request shape:
    # {"model":..., "input": "<prompt>", ...}  — input is a string, NOT messages.
    # If stream:true is ever needed, parse the SSE `chat.end` event's result.stats
    # (identical schema to the non-streaming stats block below).
    print(f"→ POST {args.lmstudio}/api/v1/chat  model={args.model}  prompt={args.prompt}")
    payload = {
        "model": args.model,
        "input": prompt_text,
        "temperature": 0.2,
        "max_output_tokens": args.max_tokens,  # native key (NOT max_tokens — that 400s)
        "stream": False,
    }
    try:
        resp, wall = api(args.lmstudio, "/api/v1/chat", 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})")

    # native response: { model_instance_id, output:[ {type:"message",content:...}, ... ],
    #                    stats:{ input_tokens, total_output_tokens, reasoning_output_tokens,
    #                            tokens_per_second, time_to_first_token_seconds, model_load_time_seconds },
    #                    response_id }
    out_items = resp.get("output", [])
    content = ""
    for it in out_items:
        if it.get("type") == "message":
            content = it.get("content", "")
            break
    stats = resp.get("stats") or {}
    tok_sec = stats.get("tokens_per_second")
    ttft = stats.get("time_to_first_token_seconds")
    comp_tokens = stats.get("total_output_tokens")
    input_tokens = stats.get("input_tokens")
    total_tokens = (input_tokens or 0) + (comp_tokens or 0) if (input_tokens or comp_tokens) else None
    # 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)
    # one-time schema dump so we can harden the extraction (write next to output)
    if not os.environ.get("GRADE_RUN_NO_SCHEMA_DUMP"):
        try:
            schema_path = os.path.join(HERE, "outputs", ".last_response_schema.json")
            slim = {k: v for k, v in resp.items() if k != "output"}
            slim["_output_types"] = [it.get("type") for it in out_items]
            json.dump(slim, open(schema_path, "w"), indent=2)
        except Exception:
            pass

    # save the raw output — extract the code block if the model wrapped it in
    # markdown fences or surrounded it with prose. Handles: (a) complete fenced
    # block, (b) prose preamble + fence, (c) bare code + trailing stray fence/prose.
    out_path = os.path.join(HERE, "outputs", f"{args.name}.{ext}")
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    import re as _re
    text = content.strip()
    pairs = list(_re.finditer(r"```(?:[a-zA-Z0-9_+-]*)?\n(.*?)```", text, _re.S))
    if pairs:
        text = pairs[0].group(1).strip()
    else:
        out, started = [], False
        for l in text.split("\n"):
            if l.strip().startswith("```"):
                break  # stray fence ends the code
            if not started and _re.match(r"^\s*(#!|import |from |async |def |class |use |pub )", l):
                started = True
            if started:
                if _re.match(r"^\s*(- |\* |\d+\. |> |# )", l) and not _re.match(r"^\s*(import |from |def |class |async |return |if |for |while |try|except|with )", l):
                    break  # trailing markdown prose
                out.append(l)
        text = "\n".join(out).strip() if out else text
    open(out_path, "w").write(text)
    print(f"✓ saved {len(text)} bytes -> {os.path.relpath(out_path, HERE)}")

    print("\n=== METRICS (from native /api/v1/chat stats) ===")
    print(f"  tok/sec      : {tok_sec}")
    print(f"  total tokens : {total_tokens}   (in={input_tokens} out={comp_tokens})")
    print(f"  TTFT (sec)   : {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": 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()