"""Many calls in flight at once, against a budget with the set-aside and one without. Python standard library only.

fire() starts all the calls together. Each call that gets past the budget waits inside the (fake) AI service until every
other call has either got in or been refused, so no call has been charged yet when the last one is checked.
"""
import threading
import time

from budget import MAX_REPLY, handle_chat

CALLS = 20  # calls started at the same moment by one user
BURST_INPUT = 20_000  # input tokens in each call
BURST_OUTPUT = 800  # output tokens each reply really has (the cap is 1,000)
USER = "user-1"


def fire(budget, calls=CALLS, estimate=BURST_INPUT, input_tokens=BURST_INPUT, output_tokens=BURST_OUTPUT):
    """Run `calls` chat calls at once. Returns (statuses, charged micro-dollars). `estimate` is the input count the server believes."""
    gate, lock = threading.Event(), threading.Lock()
    inside = [0]
    statuses = [None] * calls
    request = {"messages": [{"role": "user", "content": "a long document"}]}

    def upstream(req, max_reply):
        with lock:
            inside[0] += 1
        if not gate.wait(10):
            raise TimeoutError("the gate never opened")
        return {"text": "ok", "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}}

    def call(i):
        statuses[i] = handle_chat(budget, upstream, USER, request, 0.0, MAX_REPLY, lambda messages: estimate)[0]

    threads = [threading.Thread(target=call, args=(i,)) for i in range(calls)]
    for thread in threads:
        thread.start()
    deadline = time.monotonic() + 10
    while inside[0] + sum(s is not None for s in statuses) < calls:  # every call is either inside or already refused
        if time.monotonic() > deadline:
            gate.set()
            raise TimeoutError("calls did not settle into inside or refused")
        time.sleep(0.001)
    gate.set()
    for thread in threads:
        thread.join()
    return statuses, budget.used(USER, 0.0)