"""A REAL routing/orchestration agent (the miner's artifact source).

Orchestration layer, per query:
  1. featurize the prompt (cheap, no model call)
  2. learned logistic router -> P(cheap model suffices); route cheap vs strong
  3. VERIFY the routed answer is parseable in the task's answer format
  4. ESCALATE to the strong model if it is not

The final answer is always verbatim one pool response, so it stays grounded.
`weights` are the coefficients fit on real held-out data by fit_router.py.
"""
import json
import math
import re


def _feats(p):
    n = max(len(p), 1)
    digits = sum(c.isdigit() for c in p)
    mcq = 1.0 if re.search(r"\n\s*[A-D]\)", p) else 0.0
    return [1.0,
            math.log1p(n) / 8.0,
            10.0 * digits / n,
            mcq,
            min(p.count("?"), 3) / 3.0,
            min(len(p.split()), 200) / 200.0]


def _parseable(text, mcq):
    up = str(text).upper()
    if mcq:
        return bool(re.search(r"ANSWER[:\s]*([A-D])\b", up) or re.search(r"\b[A-D]\b", up))
    return bool(re.search(r"-?\d", str(text)))


def build_agent(weights):
    cfg = json.loads(weights.decode())
    w, cheap, strong, thr = cfg["w"], cfg["cheap"], cfg["strong"], cfg["threshold"]

    def agent(prompt, call_model):
        x = _feats(prompt)
        z = sum(wi * xi for wi, xi in zip(w, x))
        p_cheap_ok = 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, z))))
        mcq = _feats(prompt)[3] > 0
        model = cheap if p_cheap_ok >= thr else strong
        msgs = [{"role": "user", "content": prompt}]
        ans = call_model(model, msgs, {"max_tokens": 512})
        if model == cheap and not _parseable(ans, mcq):      # verify -> escalate
            ans = call_model(strong, msgs, {"max_tokens": 512})
        return ans

    return agent