"""Tool calls with nested argument schemas against a Flint SGLang server.

    pip install openai
    python3 example_tools.py --base-url http://localhost:30000/v1 [--api-key KEY]

The serve scripts start SGLang with --tool-call-parser flint, which constrains each assistant turn that
has tools to Flint's tool-call format: tool names must be ones the request lists, and the arguments of a
tool sent with "strict": true must match its JSON schema (types, enums, patterns, numeric bounds,
required fields). The tools below send "strict": true; other tools keep the format but not the schema.
Each call gets a canned result (tool_result below), which Flint's final answers report.
"""

import argparse
import json
import os

from openai import OpenAI

SAMPLING = {"temperature": 1.0, "top_p": 0.95, "extra_body": {"top_k": 20}}
MAX_TOKENS = 32768  # reasoning counts toward max_tokens
TIME = {"type": "string", "pattern": "^([01][0-9]|2[0-3]):[0-5][0-9]$"}
WEATHER = {"paris": ("sunny", 21), "tokyo": ("cloudy", 17), "lima": ("light rain", 19)}  # canned (conditions, °C)

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "create_event",
            "strict": True,
            "description": "Create a calendar event.",
            "parameters": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "date": {"type": "string", "pattern": "^[0-9]{4}-[0-9]{2}-[0-9]{2}$", "description": "YYYY-MM-DD"},
                    "start_time": {**TIME, "description": "24-hour HH:MM"},
                    "duration_minutes": {"type": "integer", "minimum": 5, "maximum": 480},
                    "attendees": {
                        "type": "array",
                        "minItems": 1,
                        "items": {
                            "type": "object",
                            "properties": {"name": {"type": "string"}, "email": {"type": "string"}},
                            "required": ["name", "email"],
                            "additionalProperties": False,
                        },
                    },
                    "reminders_minutes": {"type": "array", "items": {"type": "integer", "minimum": 0}},
                    "priority": {"type": "string", "enum": ["LOW", "NORMAL", "HIGH_PRIORITY"]},
                },
                "required": ["title", "date", "start_time", "duration_minutes", "attendees", "priority"],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "add_to_cart",
            "strict": True,
            "description": "Add products to the shopping cart.",
            "parameters": {
                "type": "object",
                "properties": {
                    "items": {
                        "type": "array",
                        "minItems": 1,
                        "items": {
                            "type": "object",
                            "properties": {
                                "sku": {"type": "string", "pattern": "^[A-Z]{3}-[0-9]{4}$"},
                                "quantity": {"type": "integer", "minimum": 1, "maximum": 99},
                                "gift_wrap": {"type": "boolean"},
                            },
                            "required": ["sku", "quantity", "gift_wrap"],
                            "additionalProperties": False,
                        },
                    },
                    "coupon": {"type": "string", "description": "Coupon code; omit when there is none"},
                },
                "required": ["items"],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "strict": True,
            "description": "Get the current weather for a city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string"},
                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
                },
                "required": ["city", "unit"],
                "additionalProperties": False,
            },
        },
    },
]

REQUESTS = [
    "Book 'Q4 planning' on October 21st, 2026 at 2:30 PM for 45 minutes with Ana (ana@example.com) and Bo "
    "(bo@example.com). Remind us 10 minutes and one hour before. It is urgent.",
    "Add two gift-wrapped mugs (SKU MUG-0042) and one notebook (SKU NTB-1001, no gift wrap) to my cart. No coupon.",
    "What's the weather in Paris, Tokyo and Lima? Use celsius.",
]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default="http://localhost:30000/v1")
    parser.add_argument("--api-key", default=os.environ.get("FLINT_API_KEY", "unused"))
    args = parser.parse_args()
    client = OpenAI(base_url=args.base_url, api_key=args.api_key)
    for request in REQUESTS:
        run(client, request)


def run(client: OpenAI, request: str) -> None:
    print(f"\n>>> {request}")
    messages = [{"role": "user", "content": request}]
    choice = client.chat.completions.create(
        model="flint", messages=messages, tools=TOOLS, max_tokens=MAX_TOKENS, **SAMPLING
    ).choices[0]
    if choice.finish_reason == "length":
        print(f"cut off: the turn used all {MAX_TOKENS} tokens (reasoning counts toward max_tokens)")
        return
    reply = choice.message
    if not reply.tool_calls:
        print("content:", reply.content)
        return
    for call in reply.tool_calls:
        print(f"{call.function.name}:", json.dumps(json.loads(call.function.arguments), ensure_ascii=False))

    # Echo the assistant turn back unchanged, answer every call, and let Flint finish the turn.
    messages.append(reply.model_dump(exclude_none=True))
    for call in reply.tool_calls:
        result = tool_result(call.function.name, json.loads(call.function.arguments))
        print(f"{call.function.name} result:", json.dumps(result, ensure_ascii=False))
        messages.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result)})
    final = client.chat.completions.create(
        model="flint", messages=messages, tools=TOOLS, max_tokens=MAX_TOKENS, **SAMPLING
    )
    print("final answer:", final.choices[0].message.content)


def tool_result(name: str, arguments: dict) -> dict:
    """Canned result for one call, standing in for a real calendar, shop and weather service."""
    if name == "get_weather":
        conditions, celsius = WEATHER.get(arguments["city"].lower(), ("clear", 20))
        temperature = celsius if arguments["unit"] == "celsius" else round(celsius * 9 / 5 + 32)
        return {
            "city": arguments["city"],
            "conditions": conditions,
            "temperature": temperature,
            "unit": arguments["unit"],
        }
    if name == "create_event":
        return {
            "status": "created",
            "event_id": "evt_1042",
            "title": arguments["title"],
            "date": arguments["date"],
            "start_time": arguments["start_time"],
        }
    return {
        "status": "added",
        "cart_id": "cart_77",
        "items_added": sum(item["quantity"] for item in arguments["items"]),
    }


if __name__ == "__main__":
    main()