"""Quick check that the install works: scores 4 sample tool calls.

    python example.py                 # answers only
    python example.py --show-prompt   # also print the exact prompt the model reads for each case
"""

import sys

from judge import ToolJudge

SHOW_PROMPT = "--show-prompt" in sys.argv

restaurant_tools = [
    {"name": "ReserveRestaurant", "description": "Make a table reservation at a specified restaurant on a given date and time for a certain number of guests."},
    {"name": "FindRestaurants", "description": "Search for restaurants based on their location, category, price range, and availability of vegetarian options and outdoor seating."},
]
booking_request = [{"role": "user", "content": "Can you book a table for four at Balthazar in New York this Friday at 8 in the evening?"}]

hotel_tools = [
    {"name": "search_hotels", "description": "Search hotels in a city by dates, area and price."},
    {"name": "book_hotel", "description": "Book a specific hotel room for given dates."},
    {"name": "play_movie", "description": "Play the selected movie."},
]
hotel_history = [   # multi turn: earlier requests and tool calls are history; only the latest user message is judged
    {"role": "user", "content": "I've rented a car in Fresno, thanks."},
    {"role": "assistant_tool_call", "content": '[{"name": "reserve_car", "arguments": {"city": "Fresno"}}]'},
    {"role": "tool", "content": '{"status": "booked", "car": "Kia Forte"}'},
    {"role": "assistant", "content": "Your Kia Forte is booked. Anything else?"},
    {"role": "user", "content": "Now help me find a hotel in Fresno for next week."},
]

cases = [
    ("1. Right tool: asked to book, agent booked", restaurant_tools, booking_request, ["ReserveRestaurant"], 1),
    ("2. Wrong tool: asked to book, agent only searched", restaurant_tools, booking_request, ["FindRestaurants"], 0),
    ("3. Tool that does not exist", restaurant_tools, booking_request, ["ReserveTable"], 0),
    ("4. Multi turn: asked for a hotel, agent played a movie", hotel_tools, hotel_history, ["play_movie"], 0),
]

print("Loading the model (first run downloads Qwen/Qwen3.5-2B)...")
judge = ToolJudge()

for title, tools, conversation, picked, expected in cases:
    print(f"\n{title}")
    if SHOW_PROMPT:
        print(judge.build_prompt(tools, conversation, picked))
    r = judge.score(tools, conversation, picked)
    verdict = "right tool" if r["answer"] == 1 else "wrong tool"
    check = "OK" if r["answer"] == expected else "UNEXPECTED"
    print(f"   agent called {picked} -> answer {r['answer']} ({verdict}), confidence {r['confidence']:.1%}, "
          f"{r['latency_ms']} ms   [{check}]")