"""Minimal OpenAI-compatible client for a Flint SGLang server, with one tool-calling round trip.

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

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

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get the current weather for a city.",
            "parameters": {
                "type": "object",
                "properties": {"city": {"type": "string", "description": "City name"}},
                "required": ["city"],
            },
        },
    }
]


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)

    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What's the weather in Paris right now?"},
    ]
    first = client.chat.completions.create(
        model="flint", messages=messages, tools=TOOLS, max_tokens=MAX_TOKENS, **SAMPLING
    )
    choice = first.choices[0]
    reply = choice.message
    print("reasoning:", getattr(reply, "reasoning_content", None))
    print("tool calls:", reply.tool_calls)
    if choice.finish_reason == "length":
        print(f"cut off: the turn used all {MAX_TOKENS} tokens (reasoning counts toward max_tokens)")
        return
    if not reply.tool_calls:
        print("content:", reply.content)
        return

    # Echo the assistant turn (including reasoning_content) back unchanged, then add tool results.
    messages.append(reply.model_dump(exclude_none=True))
    for call in reply.tool_calls:
        city = json.loads(call.function.arguments)["city"]
        messages.append(
            {"role": "tool", "tool_call_id": call.id, "content": json.dumps({"city": city, "forecast": "sunny, 21C"})}
        )
    second = client.chat.completions.create(
        model="flint", messages=messages, tools=TOOLS, max_tokens=MAX_TOKENS, **SAMPLING
    )
    print("final answer:", second.choices[0].message.content)


if __name__ == "__main__":
    main()