"""Agentic discriminator: does REAM pick the right tool with the RIGHT arguments?

The math/factual accuracy test covered the category where PPL predicted no gain.
REAM's biggest PPL gains were tool_use (27->15.6) and coding/agentic -- untested
for capability. This scores tool-call correctness right/wrong: correct function
name AND correct arguments. Argument extraction is exactly what smoothing cannot
fake -- a flatter distribution does not produce '800 by 600' or 'USD->EUR'
correctly. Same scorer for both models.

Usage: python toolcall_eval.py <model-id> <out.json>
"""
import json, re, sys, urllib.request

MODEL = sys.argv[1]
OUT = sys.argv[2] if len(sys.argv) > 2 else None

CATALOG = """Available functions (use EXACTLY these names and argument keys):
- set_timer(minutes)
- convert_currency(amount, from, to)
- book_flight(from_city, to_city, passengers)
- resize_image(width, height)
- set_thermostat(temperature)
- add_to_cart(item, quantity)
- calculate_tip(bill, percent)
- play_song(title, artist)
- get_weather(city)
- get_stock_price(ticker)
- search_products(query, max_price)
- translate(text, target_lang)
- transfer_money(amount, from_account, to_account)
- schedule_meeting(day, time)
- send_invitations(count, event)"""

SYS = (CATALOG + "\n\nCall exactly ONE function. Reply with ONLY a JSON object "
       '{"name": <function>, "arguments": {<args>}} and nothing else. '
       "Use the exact names above and the exact argument values implied by the request.")

# (request, expected_name, {arg: expected}). Numeric/string args are the
# discriminator; free-text args use keyword membership.
ITEMS = [
    ("Set a timer for 25 minutes.", "set_timer", {"minutes": "25"}),
    ("Convert 150 USD to EUR.", "convert_currency", {"amount": "150", "from": "USD", "to": "EUR"}),
    ("Book a flight from Denver to Miami for 3 passengers.", "book_flight",
     {"from_city": "Denver", "to_city": "Miami", "passengers": "3"}),
    ("Resize the image to 800 by 600.", "resize_image", {"width": "800", "height": "600"}),
    ("Set the thermostat to 68 degrees.", "set_thermostat", {"temperature": "68"}),
    ("Add 5 apples to my cart.", "add_to_cart", {"item": "apples", "quantity": "5"}),
    ("Calculate an 18% tip on an 80 dollar bill.", "calculate_tip", {"bill": "80", "percent": "18"}),
    ("Play Bohemian Rhapsody by Queen.", "play_song", {"title": "Bohemian Rhapsody", "artist": "Queen"}),
    ("What's the weather in Paris?", "get_weather", {"city": "Paris"}),
    ("Get the stock price for AAPL.", "get_stock_price", {"ticker": "AAPL"}),
    ("Find headphones under 200 dollars.", "search_products", {"query": "headphones", "max_price": "200"}),
    ("Translate 'good morning' into Spanish.", "translate", {"text": "good morning", "target_lang": "Spanish"}),
    ("Transfer 340 dollars from checking to savings.", "transfer_money",
     {"amount": "340", "from_account": "checking", "to_account": "savings"}),
    ("Schedule a meeting on Tuesday at 3pm.", "schedule_meeting", {"day": "Tuesday", "time": "3"}),
    ("Send 12 invitations to the birthday event.", "send_invitations", {"count": "12", "event": "birthday"}),
]


def ask(prompt, mt=3000):
    body = json.dumps({"model": MODEL,
                       "messages": [{"role": "system", "content": SYS},
                                    {"role": "user", "content": prompt}],
                       "temperature": 0.0, "top_p": 1.0, "top_k": 0, "min_p": 0.0,
                       "repetition_penalty": 1.0, "max_tokens": mt, "stream": False}).encode()
    r = urllib.request.Request("http://localhost:1234/v1/chat/completions", data=body,
                               headers={"Content-Type": "application/json"})
    m = json.load(urllib.request.urlopen(r, timeout=600))["choices"][0]["message"]
    return (m.get("content") or "").strip()


def parse_call(text):
    """Extract {name, arguments} from a JSON object or ```json block."""
    t = re.sub(r"```(?:json)?", "", text)
    # find the outermost {...} that has a "name"
    for m in re.finditer(r"\{.*\}", t, re.S):
        try:
            o = json.loads(m.group(0))
            if isinstance(o, dict) and "name" in o:
                return o.get("name"), (o.get("arguments") or o.get("args") or {})
        except Exception:
            continue
    return None, {}


def norm(v):
    return re.sub(r"(?<=\d)[ ,_](?=\d)", "", str(v)).strip().lower()


def check(name, args, exp_name, exp_args):
    if norm(name) != norm(exp_name):
        return False, "name"
    for k, want in exp_args.items():
        got = args.get(k)
        if got is None:
            # allow the value to appear under any key (arg-naming can differ)
            if any(norm(want) == norm(val) or norm(want) in norm(val)
                   for val in args.values()):
                continue
            return False, f"missing {k}={want}"
        if norm(want) == norm(got) or norm(want) in norm(got):
            continue
        return False, f"{k}: want {want!r} got {got!r}"
    return True, "ok"


def self_test():
    n, a = parse_call('```json\n{"name": "set_timer", "arguments": {"minutes": 25}}\n```')
    assert n == "set_timer" and str(a["minutes"]) == "25"
    ok, _ = check(n, a, "set_timer", {"minutes": "25"}); assert ok
    ok, why = check("set_timer", {"minutes": "30"}, "set_timer", {"minutes": "25"})
    assert not ok, why
    ok, _ = check("resize_image", {"w": "800", "h": "600"}, "resize_image",
                  {"width": "800", "height": "600"})
    assert ok, "value-under-any-key fallback"
    print("[selftest] tool-call parser ok", flush=True)


if __name__ == "__main__":
    self_test()
    if MODEL == "SELFTEST":
        sys.exit(0)
    ok_n = 0; results = []
    for req, en, ea in ITEMS:
        try:
            raw = ask(req)
        except Exception as e:
            print(f"  ERROR {e}", flush=True); continue
        n, a = parse_call(raw)
        ok, why = check(n, a, en, ea)
        ok_n += ok
        results.append({"req": req, "ok": ok, "why": why, "name": n, "args": a})
        print(f"  [{'OK ' if ok else 'XX '}] want {en}({ea}) -> got {n}({a})  {'' if ok else why}", flush=True)
    print(f"\n{MODEL}  TOOL-CALL {ok_n}/{len(ITEMS)}  ({100*ok_n/len(ITEMS):.0f}%)", flush=True)
    if OUT:
        json.dump({"model": MODEL, "score": [ok_n, len(ITEMS)], "results": results},
                  open(OUT, "w"), indent=2)