"""A budget for each signed-in user, counted in money and checked before each AI call. Python standard library only.

Prices are the example prices from lesson 7 ($2 and $10 per million tokens), not a quote. Money is counted in
micro-dollars (millionths of a dollar) as whole numbers, so $2 per million tokens is exactly 2 micro-dollars per token.

A call goes through three steps. reserve() sets aside the most the call can cost and refuses if that does not fit.
settle() charges what the call really cost and gives the rest back. release() gives everything back when the service answered with an error.
Time is a parameter (seconds) so a test can move it without waiting.
"""
import math
import threading
from dataclasses import dataclass

PRICE_IN = 2  # micro-dollars per input token
PRICE_OUT = 10  # micro-dollars per output token
MAX_REPLY = 1000  # the reply cap the server sends with every call (lesson 9's max_tokens)
DAY = 86_400  # seconds
MESSAGE_OVERHEAD = 10  # tokens added per message by the chat format: an assumption; check your provider's counter


def cost(input_tokens, output_tokens):
    """Micro-dollars a call costs."""
    return input_tokens * PRICE_IN + output_tokens * PRICE_OUT


def dollars(micro):
    """Text such as "$0.25". Whole cents, halves round up (widget.html does the same sum)."""
    cents = (micro + 5_000) // 10_000
    return f"${cents // 100}.{cents % 100:02d}"


def how_long(seconds):
    """Text such as "9 h 20 min", rounded up to a whole minute."""
    hours, minutes = divmod(max(1, math.ceil(seconds / 60)), 60)
    return f"{hours} h {minutes} min" if hours else f"{minutes} min"


def input_tokens_upper_bound(messages):
    """A count that never falls short for the text in `messages`: a token covers at least one byte of text, so bytes >= tokens (plus the per-message extra). The caller adds the system prompt and tool definitions."""
    return sum(len(m["content"].encode("utf-8")) + MESSAGE_OVERHEAD for m in messages)


class OverBudget(Exception):
    """reserve() refused. `left` is what is still free, `wait` the seconds until the call fits (None: it never will)."""

    def __init__(self, amount, left, wait):
        super().__init__(f"needs {amount}, {left} left")
        self.amount, self.left, self.wait = amount, left, wait


@dataclass(frozen=True)
class Hold:
    user: str
    amount: int  # micro-dollars set aside
    period: int  # the day it was set aside in (steady refill: always 0)


class Budget:
    """Each user may use `limit` micro-dollars. refill="daily": the count restarts at 00:00 UTC.
    refill="steady": the count drains by limit/86400 per second, so the whole budget is back after a day of no use."""

    def __init__(self, limit, refill="daily"):
        if refill not in ("daily", "steady"):
            raise ValueError("refill is daily or steady")
        self.limit, self.refill = limit, refill
        self._state = {}  # user -> (used, stamp); stamp is the day number (daily) or the time of the last update (steady)
        self._lock = threading.Lock()  # one process only: with two processes the state must live in one shared store

    def _fresh(self, user, now):
        """(used, stamp) for this user brought up to `now`. Call it with the lock held."""
        used, stamp = self._state.get(user, (0, None))
        if self.refill == "daily":
            day = int(now // DAY)
            return (used if stamp == day else 0), day
        if stamp is not None:
            used = max(0, used - (now - stamp) * self.limit / DAY)
        return used, now

    def used(self, user, now):
        with self._lock:
            return self._fresh(user, now)[0]

    def reserve(self, user, amount, now):
        """Set `amount` aside, or raise OverBudget. The check and the set-aside happen under one lock, so two calls cannot both take the last of it."""
        with self._lock:
            used, stamp = self._fresh(user, now)
            if used + amount > self.limit:
                excess = used + amount - self.limit
                if amount > self.limit:
                    wait = None  # bigger than the whole budget: waiting never helps
                elif self.refill == "daily":
                    wait = DAY - now % DAY
                else:
                    wait = excess * DAY / self.limit
                raise OverBudget(amount, max(0, self.limit - used), wait)
            self._state[user] = (used + amount, stamp)
            return Hold(user, amount, stamp if self.refill == "daily" else 0)

    def settle(self, hold, actual, now):
        """Replace the set-aside with the real cost. If the real cost is higher, the user is charged it all and may go over."""
        with self._lock:
            used, stamp = self._fresh(hold.user, now)
            if self.refill == "daily" and hold.period != stamp:
                return  # the day this call was set aside in has ended, and its count with it
            self._state[hold.user] = (max(0, used + actual - hold.amount), stamp)

    def release(self, hold, now):
        """The service answered with an error, or never got the call: give the whole set-aside back."""
        self.settle(hold, 0, now)


class CheckOnlyBudget(Budget):
    """The first version most people write: look at what is already charged, make the call, charge afterwards. Nothing is set aside."""

    def reserve(self, user, amount, now):
        return super().reserve(user, 0, now)


def refusal(refused, limit, now):
    """The JSON body for a refused call: a code the app can test, a message a person can read, and when the budget is back."""
    if refused.wait is None:
        return {"error": "request_too_big",
                "message": f"This request could cost up to {dollars(refused.amount)}, more than your whole {dollars(limit)} budget. Send less text."}
    return {"error": "budget_used_up", "resets_at": int(now + refused.wait),
            "message": f"Your AI budget has {dollars(refused.left)} left and this call needs up to {dollars(refused.amount)}. "
                       f"It refills in {how_long(refused.wait)}."}


def handle_chat(budget, upstream, user, request, now, max_reply=MAX_REPLY, count_tokens=input_tokens_upper_bound):
    """Run one chat call for a signed-in user. Returns (status, body).

    `user` is the account id from your sign-in check, never a value the browser sends. `upstream(request, max_reply)`
    calls the AI service and returns {"text": ..., "usage": {"input_tokens": ..., "output_tokens": ...}}. On a
    timeout it must raise Python's own TimeoutError, so map your HTTP library's timeout error to it.
    """
    if user is None:
        return 401, {"error": "sign_in", "message": "Sign in to use the assistant."}
    amount = cost(count_tokens(request["messages"]), max_reply)  # the most this call can cost
    try:
        hold = budget.reserve(user, amount, now)
    except OverBudget as refused:
        return 429, refusal(refused, budget.limit, now)  # no short Retry-After: it tells clients to retry soon, and the budget is back in hours
    try:
        answer = upstream(request, max_reply)
    except TimeoutError:
        # the provider may still bill a call that timed out, so keep the set-aside: the safe side.
        # It stays counted until the daily reset, or drains with a steady refill.
        return 502, {"error": "upstream_timeout", "message": "The AI service did not answer in time. The call may still count against your budget."}
    except Exception:
        budget.release(hold, now)  # the service answered with an error, so give the set-aside back
        return 502, {"error": "upstream_failed", "message": "The AI service failed. Try again in a moment."}
    usage = answer["usage"]
    actual = cost(usage["input_tokens"], usage["output_tokens"])
    budget.settle(hold, actual, now)
    return 200, {"text": answer["text"], "charged": actual}