"""Jedi runtime: answer with the cheapest agent, let the 0.5B verifier check it, escalate up the ladder if needed.

Usage:
  JEDI_ENDPOINTS='{"glm-5.3-flash": "http://host:8000/v1", ...}' python jedi.py [--threshold 0.7] [--code] "task"

The threshold is the single cost-vs-quality dial: raise it for accuracy (more escalations), lower it for savings.
Prints the final answer; the per-rung trace goes to stderr as JSON lines like
  {"agent": "glm-5.3-flash", "verdict": "escalate", "confidence": 0.87, "escalate_to": "qwen3.8-27b"}
"""
import argparse, json, math, os, re, sys
import httpx
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from common import CODE_SYSTEM, LADDER, POOL, JEDI_SYSTEM, SOLVER_SYSTEM, endpoints, est_tokens, judge_view, max_out

MODEL = os.environ.get("JEDI_MODEL", "podstack/Jedi")


class Jedi:
    def __init__(self, path=MODEL):
        self.tok = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).eval()
        if torch.cuda.is_available():
            self.model.cuda()
        try:
            from huggingface_hub import hf_hub_download
            cal_file = os.path.join(path, "calibration.json") if os.path.isdir(path) else \
                hf_hub_download(path, "calibration.json")
            cal = json.load(open(cal_file))
        except Exception:  # noqa: BLE001
            cal = {"a": 1.0, "b": 0.0}
        self.a, self.b = cal["a"], cal["b"]
        self.ids = [self.tok.encode(w, add_special_tokens=False)[0] for w in ("accept", "escalate")]

    def p_correct(self, task, answer, agent, kind, finish=None):
        """Calibrated probability that `answer` (written by `agent`) is correct."""
        enc = self.tok.apply_chat_template(
            [{"role": "system", "content": JEDI_SYSTEM},
             {"role": "user", "content": judge_view(task, answer, agent, kind, finish)}],
            add_generation_prompt=True, return_tensors="pt", return_dict=True).to(self.model.device)
        with torch.no_grad():
            logits = self.model(**enc).logits[0, -1].float()
        gap = (logits[self.ids[0]] - logits[self.ids[1]]).item()
        return 1 / (1 + math.exp(-(self.a * gap + self.b)))


def ask(agent, task, code):
    base = endpoints()[agent].rstrip("/")
    msgs = [{"role": "system", "content": CODE_SYSTEM if code else SOLVER_SYSTEM}, {"role": "user", "content": task}]
    r = httpx.post(f"{base}/chat/completions", timeout=1800, json={
        "model": POOL[agent]["model"], "messages": msgs, "max_tokens": max_out(agent, est_tokens(task))})
    r.raise_for_status()
    c = r.json()["choices"][0]
    return c["message"].get("content") or "", c.get("finish_reason")


def infer_kind(task, code):
    """Task type shown to Jedi (its training data used: code, mc, num, math, exact)."""
    if code:
        return "code"
    return "mc" if re.search(r"^\(?[A-J][.)]\s", task, flags=re.M) else "math"


def solve(task, threshold=0.7, code=False, kind=None, judge=None):
    """Run the ladder. Returns (final answer text, trace)."""
    judge = judge or Jedi()
    kind = kind or infer_kind(task, code)
    rungs = [a for a in LADDER if a in endpoints()]
    trace, text = [], ""
    for i, agent in enumerate(rungs):
        try:
            text, finish = ask(agent, task, code)
        except httpx.HTTPError as e:
            if i == 0:
                raise
            trace.append({"agent": agent, "verdict": "unavailable", "error": repr(e)[:120]})
            return text, trace  # keep the best answer so far
        if i == len(rungs) - 1:
            trace.append({"agent": agent, "verdict": "final", "confidence": None, "escalate_to": None})
            return text, trace
        p = judge.p_correct(task, text, agent, kind, finish)
        verdict = "accept" if p >= threshold else "escalate"
        trace.append({"agent": agent, "verdict": verdict, "confidence": round(p if verdict == "accept" else 1 - p, 3),
                      "escalate_to": None if verdict == "accept" else rungs[i + 1]})
        if verdict == "accept":
            return text, trace


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("task", nargs="+")
    ap.add_argument("--threshold", type=float, default=0.7, help="min P(correct) to accept an answer")
    ap.add_argument("--code", action="store_true", help="ask agents for a Python code solution")
    ap.add_argument("--kind", choices=["math", "num", "mc", "code", "exact"], help="task type (default: inferred)")
    args = ap.parse_args()
    answer, trace = solve(" ".join(args.task), args.threshold, args.code, args.kind)
    for step in trace:
        print(json.dumps(step), file=sys.stderr)
    print(answer)