"""Shared config for Jedi: the agent pool, the escalation ladder, prompts, capabilities, graders and the judge view."""
import json, os, re

# agent -> served model id, max concurrent requests we send, relative cost per token
POOL = {
    "gemma-3-1b":         dict(model="google/gemma-3-1b-it", concurrency=32, cost=0.02),
    "kimi-vl":            dict(model="Kimi-VL-A3B-Thinking", concurrency=64, cost=0.05),
    "nemotron-lightning": dict(model="nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", concurrency=64, cost=0.10),
    "glm-5.3-flash":      dict(model="zai-org/GLM-5.3-Flash", concurrency=32, cost=0.20),
    "qwen3.8-27b":        dict(model="Qwen/Qwen3.8-27B", concurrency=12, cost=0.50),
    "nemotron-ultra":     dict(model="nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", concurrency=64, cost=1.00),
}
AGENTS = list(POOL)  # ordered cheapest -> most expensive
STRONGEST = "nemotron-ultra"
CONTEXT = {"gemma-3-1b": 32768, "kimi-vl": 8192, "nemotron-lightning": 131072, "glm-5.3-flash": 32768,
           "qwen3.8-27b": 32768, "nemotron-ultra": 262144}
VISION = {"kimi-vl"}  # only agent served with image input enabled
MAX_OUT = 8192


LADDER = ["glm-5.3-flash", "qwen3.8-27b", "nemotron-ultra"]  # cheap -> mid -> strongest; Jedi checks each rung


def endpoints() -> dict:
    """Agent -> OpenAI-compatible base URL (".../v1"), from $JEDI_ENDPOINTS: inline JSON or a path to a JSON file."""
    raw = os.environ["JEDI_ENDPOINTS"]
    return json.load(open(raw)) if os.path.exists(raw) else json.loads(raw)

SOLVER_SYSTEM = ("Solve the task. Be concise. End your response with a final line exactly of the form:\n"
                 "Final answer: <answer>\nFor multiple-choice questions, the answer is the option letter only.")
CODE_SYSTEM = ("Write a correct, self-contained Python solution. Put the complete code (imports, required function "
               "names and signatures) in a single ```python code block. Do not include tests or example usage.")

JEDI_SYSTEM = (
    "You are Jedi, a verifier. You are shown a task and a candidate answer written by an AI agent. Do not solve "
    "the task yourself. Judge whether the candidate's final answer is correct: look for reasoning that does not hold "
    "up, arithmetic that does not add up, code that would fail its tests, hedging, truncation or a wrong answer "
    "format. Reply with one word: accept if the answer is correct, escalate if it is likely wrong."
)
VERDICTS = ("accept", "escalate")


def est_tokens(text: str) -> int:
    """Conservative token estimate (real tokenizers average ~3.5-4 chars/token on English)."""
    return len(text) // 3 + 1


def max_out(agent: str, prompt_tokens: int) -> int:
    return min(MAX_OUT, CONTEXT[agent] - prompt_tokens - 256)


def capable(agent: str, prompt: str, has_image: bool) -> bool:
    if has_image and agent not in VISION:
        return False
    return max_out(agent, est_tokens(prompt) + 1500 * has_image) >= 1024


def truncate(prompt: str, head: int = 5000, tail: int = 1500) -> str:
    return prompt if len(prompt) <= head + tail + 7 else prompt[:head] + "\n[...]\n" + prompt[-tail:]


def judge_view(prompt: str, answer: str, agent: str, kind: str, finish: str | None) -> str:
    """What Jedi reads: who answered and how the answer ended, then the task and the visible answer (head + tail)."""
    answer = strip_think(answer)
    head = f"[agent: {agent}; task type: {kind}; answer ~{est_tokens(answer)} tokens" + \
           ("; answer was cut off at the token limit" if finish == "length" else "") + "]"
    return f"{head}\n### Task\n{truncate(prompt, 2500, 1000)}\n\n### Candidate answer\n{truncate(answer, 1500, 2500)}"


def strip_think(text: str) -> str:
    text = text or ""
    if "</think>" in text:
        text = text.rsplit("</think>", 1)[1]
    text = re.sub(r"◁think▷.*?◁/think▷", "", text, flags=re.S)  # Kimi-VL thinking markers
    return text.strip()


def extract_final(text: str) -> str:
    text = strip_think(text)
    m = re.findall(r"final answer\s*[:：]\s*(.+)", text, flags=re.I)
    if m:
        return m[-1].strip()
    m = re.findall(r"^\s*answer\s*[:：]\s*(.+)", text, flags=re.I | re.M)
    if m:
        return m[-1].strip()
    boxed = last_boxed(text)
    if boxed is not None:
        return boxed
    lines = [l for l in text.splitlines() if l.strip()]
    return lines[-1].strip() if lines else ""


def extract_code(text: str) -> str:
    text = strip_think(text)
    blocks = re.findall(r"```(?:python|py)?\s*\n(.*?)```", text, flags=re.S)
    return blocks[-1] if blocks else ""


def last_boxed(s: str):
    i = s.rfind("\\boxed{")
    if i < 0:
        return None
    i += len("\\boxed{")
    depth, j = 1, i
    while j < len(s) and depth:
        depth += {"{": 1, "}": -1}.get(s[j], 0)
        j += 1
    return s[i:j - 1] if depth == 0 else None


def _norm_math(s: str) -> str:
    s = s.strip().strip("$").strip()
    b = last_boxed(s)
    if b is not None:
        s = b
    for a, b2 in [("\\left", ""), ("\\right", ""), ("\\!", ""), ("\\,", ""), ("\\;", ""), ("dfrac", "frac"),
                  ("tfrac", "frac"), ("^\\circ", ""), ("^{\\circ}", ""), ("\\%", ""), ("%", ""), ("\\text{", "{"),
                  ("\\mbox{", "{"), (" ", ""), ("\\$", "")]:
        s = s.replace(a, b2)
    s = re.sub(r"\\frac(\d)(\d)", r"\\frac{\1}{\2}", s)
    s = s.rstrip(".")
    if s.startswith("{") and s.endswith("}") and s.count("{") == 1:
        s = s[1:-1]
    return s


def _num(s: str):
    s = s.replace(",", "").replace("$", "").strip()
    m = re.findall(r"-?\d+(?:\.\d+)?", s)
    return float(m[-1]) if m else None


def grade(kind: str, pred: str, gold: str) -> bool:
    """Grades extracted answers. Code is graded by execution (see gen.run_tests), not here."""
    if not pred:
        return False
    if kind == "mc":
        m = re.search(r"\b([A-J])\b", pred.upper().replace("(", " ").replace(")", " "))
        return bool(m) and m.group(1) == gold.upper()
    if kind == "num":
        p, g = _num(pred), _num(gold)
        return p is not None and g is not None and abs(p - g) < 1e-6 * max(1, abs(g))
    if kind == "relaxed":  # ChartQA convention: numbers within 5%, otherwise case-insensitive exact match
        pred = last_boxed(pred) or pred
        pred = re.sub(r"\\(?:mathbf|textbf|text|mathrm)\{([^}]*)\}", r"\1", pred).replace("\\(", "").replace("\\)", "")
        p, g = _num(pred), _num(gold)
        if g is not None and p is not None and re.fullmatch(r"[-\d.,%$ ]+", gold.strip()):
            return abs(p - g) <= 0.05 * abs(g) if g else p == 0
        return pred.strip().strip(".").lower() == gold.strip().lower()
    if kind == "math":
        p, g = _norm_math(pred), _norm_math(gold)
        if p == g:
            return True
        try:
            return abs(float(p) - float(g)) < 1e-6
        except ValueError:
            return False
    if kind == "exact":  # BIG-Bench Hard targets: an option like "(B)", or a short string / number / list
        norm = lambda x: re.sub(r"\s+", " ", x.strip().strip("*").strip().rstrip(".").lower())
        p, g = norm(pred), norm(gold)
        if re.fullmatch(r"\([a-r]\)", g):
            m = re.match(r"^\(?([a-r])\)?(?:\W|$)", p)
            return bool(m) and m.group(1) == g[1]
        return p == g
    raise ValueError(kind)