"""two_budgets.py - two requests set an agent's budget. Which one stops the agent?

    python two_budgets.py

Paperclip's API pages give PATCH /api/agents/{id} with budgetMonthlyCents.
The server also has PATCH /api/agents/{id}/budgets. This check sets a budget
of 100 cents each way, on two agents, and then reports 120 cents of cost.
"""
import time

import scratch

WAYS = [("PATCH /agents/{id}", ""), ("PATCH /agents/{id}/budgets", "/budgets")]


def main():
    made, boss = scratch.company("Check two budgets")
    for label, ending in WAYS:
        agent = scratch.idle_agent(made, boss, "Agent" + (ending or "/plain"))
        scratch.send("PATCH", "/agents/" + agent["id"] + ending, {"budgetMonthlyCents": 100})
        for _ in range(3):
            now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
            scratch.send("POST", "/companies/" + made["id"] + "/cost-events",
                         {"agentId": agent["id"], "provider": "pretend", "model": "pretend",
                          "inputTokens": 1000, "outputTokens": 200, "costCents": 40,
                          "occurredAt": now})
        time.sleep(2)
        agent = scratch.send("GET", "/agents/" + agent["id"])
        print(f"{label}: budget {agent['budgetMonthlyCents']}, spent "
              f"{agent['spentMonthlyCents']}, status {agent['status']}, "
              f"pause reason {agent['pauseReason']}")
    scratch.remove(made)


if __name__ == "__main__":
    main()