import json
import os
from openai import OpenAI


STOCK = {"QR-001": 7, "QR-002": 0}
TOOL = {"type": "function", "function": {
    "name": "get_stock", "description": "查询演示库存；没有写入或下单操作。",
    "parameters": {"type": "object", "properties": {"sku": {"type": "string"}},
                   "required": ["sku"], "additionalProperties": False}}}


def execute(name, arguments):
    if name != "get_stock":
        return {"error": "tool_not_allowed"}
    try:
        args = json.loads(arguments)
        if not isinstance(args, dict) or set(args) != {"sku"} or not isinstance(args["sku"], str):
            return {"error": "invalid_arguments"}
        return {"sku": args["sku"], "stock": STOCK.get(args["sku"]), "source": "local_demo_fixture"}
    except (json.JSONDecodeError, TypeError):
        return {"error": "invalid_json"}


def main():
    client = OpenAI(api_key=os.environ["QUICKROUTER_API_KEY"],
                    base_url="https://api.quickrouter.ai/v1", timeout=60, max_retries=0)
    messages = [{"role": "user", "content": "调用 get_stock 查询 QR-001 库存，再用中文回答。"}]
    ledger = []
    for turn in range(4):
        result = client.chat.completions.create(model=os.environ["QUICKROUTER_MODEL"],
                                                messages=messages, tools=[TOOL])
        ledger.append({"turn": turn + 1, "usage": result.usage.model_dump() if result.usage else None})
        print("usage turn:", json.dumps(ledger[-1]))
        choice = result.choices[0]
        message = choice.message
        if message.tool_calls:
            if choice.finish_reason != "tool_calls":
                raise RuntimeError("Tool call output is incomplete")
            messages.append(message.model_dump(exclude_none=True))
            for call in message.tool_calls:
                if call.type != "function":
                    raise RuntimeError("Unsupported tool type")
                output = execute(call.function.name, call.function.arguments)
                print("tool:", json.dumps(output, ensure_ascii=False))
                messages.append({"role": "tool", "tool_call_id": call.id,
                                 "content": json.dumps(output, ensure_ascii=False)})
        else:
            if choice.finish_reason != "stop" or not message.content:
                raise RuntimeError("Final answer is incomplete")
            print(message.content)
            print("usage ledger:", json.dumps(ledger))
            return
    raise RuntimeError("Tool loop exceeded four model calls")


if __name__ == "__main__":
    main()