"""Tool-calling battery for candidate brain models on a scratch llama port.
Usage: battery.py <base_url>   e.g. python battery.py http://127.0.0.1:8082/v1
Scores: tool call correctness + grounded final answer + timing, per question.
"""
import sys, json, time
sys.path.insert(0, r"C:\oracle")
import server

server.BRAIN_URL = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:8082/v1/chat/completions"
server.BRAIN_TIMEOUT = 120

QUESTIONS = [
    ("date",      "What is today's date?"),
    ("weather",   "What is the weather in Oceanside, California right now?"),
    ("math",      "what is 137 times 89?"),
    ("convert",   "how many miles is 42 kilometers?"),
    ("web",       "Who won the most recent Super Bowl? Use web search."),
    ("plain",     "In two short sentences, why is the sky blue?"),
]

def agent_loop(question):
    msgs = [{"role": "system", "content": server.PERSONA},
            {"role": "user", "content": question}]
    calls = []
    for _pass in range(2):
        msg = server.llm_respond(msgs, server.TOOLS)
        content = (msg.get("content") or "").strip()
        tcs = msg.get("tool_calls")
        if not tcs:
            return calls, content or "(empty)"
        for i, tc in enumerate(tcs):
            fn = (tc.get("function") or {}).get("name", "?")
            try: args = json.loads((tc.get("function") or {}).get("arguments") or "{}")
            except Exception: args = {}
            calls.append((fn, args))
        msgs.append({"role": "assistant", "content": content or None, "tool_calls": tcs})
        for i, tc in enumerate(tcs):
            fn = (tc.get("function") or {}).get("name", "?")
            try: args = json.loads((tc.get("function") or {}).get("arguments") or "{}")
            except Exception: args = {}
            res = server.run_tool(fn, args)
            msgs.append({"role": "tool", "tool_call_id": tc.get("id") or f"call_{i}", "content": res})
    return calls, "(loop exhausted)"

print(f"=== TOOL BATTERY against {server.BRAIN_URL} ===")
for tag, q in QUESTIONS:
    t0 = time.time()
    try:
        calls, ans = agent_loop(q)
        dt = time.time() - t0
        calls_s = "; ".join(f"{fn}({json.dumps(a)[:60]})" for fn, a in calls) or "NONE"
        print(f"\n[{tag}] {dt:5.1f}s  tools: {calls_s}")
        print(f"    ans: {ans[:180]}")
    except Exception as e:
        print(f"\n[{tag}] ERROR {type(e).__name__}: {e}")