#!/usr/bin/env python3
"""Eyeball BTL-4 Compact's tool use before running the full BFCL gate.

Ten prompts covering the five behaviours BTL-3 Compact was scored on: a single
call, picking the right tool from several, two calls in parallel, two *different*
tools in parallel, and knowing when to make no call at all. Parallel-multiple is
the one to watch -- it was BTL-3 Compact's weakest category at 3/10.

Shares the system prompt and parser with bfcl_compact.py so what you see here is
what the benchmark will score.

  python probe_tools.py                    # all ten
  python probe_tools.py --ask "your question here"
"""
from __future__ import annotations

import argparse
import json

from bfcl_compact import REPO, FILENAME, SYS, parse_tool_calls

TOOLS = [
    {"name": "get_weather",
     "description": "Get the current weather for a city.",
     "parameters": {"type": "object", "properties": {
         "city": {"type": "string", "description": "City name"},
         "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}},
         "required": ["city"]}},
    {"name": "convert_currency",
     "description": "Convert an amount between two currencies.",
     "parameters": {"type": "object", "properties": {
         "amount": {"type": "number"},
         "from_currency": {"type": "string", "description": "ISO code, e.g. USD"},
         "to_currency": {"type": "string", "description": "ISO code, e.g. EUR"}},
         "required": ["amount", "from_currency", "to_currency"]}},
    {"name": "search_flights",
     "description": "Search available flights between two airports on a date.",
     "parameters": {"type": "object", "properties": {
         "origin": {"type": "string"}, "destination": {"type": "string"},
         "date": {"type": "string", "description": "YYYY-MM-DD"}},
         "required": ["origin", "destination", "date"]}},
    {"name": "send_email",
     "description": "Send an email.",
     "parameters": {"type": "object", "properties": {
         "to": {"type": "string"}, "subject": {"type": "string"},
         "body": {"type": "string"}},
         "required": ["to", "subject", "body"]}},
    {"name": "stock_price",
     "description": "Get the latest share price for a ticker symbol.",
     "parameters": {"type": "object", "properties": {
         "ticker": {"type": "string"}},
         "required": ["ticker"]}},
]

# (prompt, what a correct model should do) -- the expectation is for your eyes,
# nothing here is auto-scored.
PROBES = [
    ("What's the weather in Lagos?",
     "single: get_weather(city='Lagos')"),
    ("How much is 250 US dollars in Japanese yen?",
     "single, right tool from five: convert_currency"),
    ("What's the weather in Lagos and in Tokyo?",
     "parallel: get_weather twice"),
    ("Give me the weather in Berlin and the share price of NVDA.",
     "parallel-multiple: two DIFFERENT tools"),
    ("Convert 100 GBP to EUR and 100 GBP to USD, and tell me Tesla's stock price.",
     "parallel-multiple: three calls, two tools"),
    ("Find me flights from LHR to CDG on 2026-09-14.",
     "single with a date argument"),
    ("Write me a haiku about the rain.",
     "ABSTAIN: no tool applies"),
    ("What do you think is the best programming language?",
     "ABSTAIN: opinion, no tool"),
    ("Email ada@example.com with the subject 'Q3 numbers' saying the figures are approved.",
     "single with three string args"),
    ("What's the weather in Paris, and email it to sam@example.com with subject 'Paris'?",
     "parallel-multiple: get_weather + send_email"),
]


def run(llm, question: str, expect: str | None = None) -> None:
    msgs = [{"role": "system", "content": SYS + json.dumps(TOOLS)},
            {"role": "user", "content": question}]
    out = llm.create_chat_completion(messages=msgs, max_tokens=512, temperature=0.0)
    raw = out["choices"][0]["message"].get("content") or ""
    calls = parse_tool_calls(raw)

    print(f"\n\033[1m❯ {question}\033[0m")
    if expect:
        print(f"  \033[2mexpect: {expect}\033[0m")
    if calls:
        for c in calls:
            args = ", ".join(f"{k}={v!r}" for k, v in c["arguments"].items())
            print(f"  \033[32m→ {c['name']}({args})\033[0m")
    else:
        body = " ".join(raw.split())[:200]
        print(f"  \033[33m→ no tool call\033[0m  {body}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ask", help="run a single custom question")
    ap.add_argument("--model", default=None, help="local .gguf path")
    ap.add_argument("--ctx", type=int, default=8192)
    args = ap.parse_args()

    from llama_cpp import Llama
    path = args.model
    if path is None:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(repo_id=REPO, filename=FILENAME)

    print("loading onto the GPU ...", flush=True)
    llm = Llama(model_path=path, n_gpu_layers=-1, n_ctx=args.ctx, verbose=False)

    if args.ask:
        run(llm, args.ask)
        return
    for q, expect in PROBES:
        run(llm, q, expect)
    print("\n\033[2mparallel-multiple is the one that matters: BTL-3 Compact "
          "scored 3/10 there.\033[0m")


if __name__ == "__main__":
    main()