"""One day of AI calls from ten users, one of them heavy, under three rules. Python standard library only.

A toy, not a model of any real app. Ordinary users chat in a few short sessions; the heavy user sends long documents
all day, one call every few minutes. The three rules differ only in what the server checks before it makes a call:
nothing, a number of calls per minute for each user, or a daily budget in money for each user. A refused call is lost
(the user is told and does not retry). The same seed always gives the same day; widget.html has the same code.
"""
from dataclasses import dataclass

from budget import DAY, MAX_REPLY, Budget, OverBudget, cost

USERS = 10  # user 0 is the heavy user, users 1 to 9 are ordinary
HEAVY = 0
SEED = 7  # the day the lesson shows
RPM = 5  # calls per minute each user may make under the "requests per minute" rule
BUDGET = 250_000  # micro-dollars each user may spend in a day under the budget rule ($0.25)
HEAVY_FROM, HEAVY_TO = 6 * 3600, 22 * 3600  # the heavy user works from 06:00 to 22:00
RULES = (("none", "No limit"), ("rpm", "Requests per minute"), ("budget", "Daily cost budget"))
START_RULE = "rpm"  # the widget opens on this rule

# What one call is made of, as (lowest, highest) token counts. Ordinary users: 3 to 5 sessions a day of 4 to 8 calls, 20 to 89 s apart.
ORDINARY_IN, ORDINARY_OUT = (300, 900), (80, 280)
HEAVY_IN, HEAVY_OUT = (16_000, 24_000), (700, 1000)
HEAVY_GAP = (240, 360)  # seconds between the heavy user's calls
SESSIONS, SESSION_CALLS, CALL_GAP = (3, 5), (4, 8), (20, 89)
LONGEST_SESSION = 8 * 89  # seconds; a session starts early enough to end inside the day

MASK = 0xFFFFFFFF


def random_numbers(seed):
    """Endless numbers in [0, 1), the same on every machine (the mulberry32 generator, as in lesson 10; widget.html has the same code)."""
    state = seed & MASK
    while True:
        state = (state + 0x6D2B79F5) & MASK
        t = ((state ^ (state >> 15)) * (state | 1)) & MASK
        t ^= (t + (((t ^ (t >> 7)) * (t | 61)) & MASK)) & MASK
        yield ((t ^ (t >> 14)) & MASK) / 4294967296


class TokenBucket:
    """From lesson 9: holds up to `capacity` tickets and gains `rate` tickets per second. A call spends one."""

    def __init__(self, capacity, rate, now):
        self.capacity, self.rate = capacity, rate
        self.level, self.last = float(capacity), now

    def take(self, now):
        """Spend one ticket if there is one. Return True if the call is allowed."""
        self.level = min(self.capacity, self.level + (now - self.last) * self.rate)
        self.last = now
        if self.level >= 1:
            self.level -= 1
            return True
        return False


def between(draw, low, high):
    """A whole number from low to high, both included."""
    return low + int(draw() * (high - low + 1))


def make_day(seed):
    """Every call of the day as (second, user, input tokens, output tokens), in time order."""
    draw = random_numbers(seed).__next__
    calls = []
    t = HEAVY_FROM
    while t < HEAVY_TO:
        calls.append((t, HEAVY, between(draw, *HEAVY_IN), between(draw, *HEAVY_OUT)))
        t += between(draw, *HEAVY_GAP)
    for user in range(1, USERS):
        for _ in range(between(draw, *SESSIONS)):
            t = between(draw, 0, DAY - LONGEST_SESSION)
            for _ in range(between(draw, *SESSION_CALLS)):
                calls.append((t, user, between(draw, *ORDINARY_IN), between(draw, *ORDINARY_OUT)))
                t += between(draw, *CALL_GAP)
    return sorted(calls, key=lambda c: (c[0], c[1]))  # equal times keep the order they were made in


@dataclass(frozen=True)
class Day:
    rule: str
    served: tuple  # per user: calls answered
    refused: tuple  # per user: calls refused
    spent: tuple  # per user: micro-dollars charged

    @property
    def bill(self):
        return sum(self.spent)

    def percent(self, user):
        """The user's share of the bill in whole percent, halves round up (whole-number arithmetic, so JavaScript agrees)."""
        return (200 * self.spent[user] + self.bill) // (2 * self.bill) if self.bill else 0

    @property
    def ordinary_refused(self):
        return sum(self.refused[1:])

    @property
    def ordinary_asked(self):
        return sum(self.served[1:]) + self.ordinary_refused


def simulate_day(rule, seed=SEED, rpm=RPM, budget=BUDGET):
    served, refused, spent = [0] * USERS, [0] * USERS, [0] * USERS
    buckets = [TokenBucket(rpm, rpm / 60, 0) for _ in range(USERS)]
    wallet = Budget(budget)
    for t, user, tokens_in, tokens_out in make_day(seed):
        allowed = True
        if rule == "rpm":
            allowed = buckets[user].take(t)
        elif rule == "budget":
            try:
                wallet.settle(wallet.reserve(user, cost(tokens_in, MAX_REPLY), t), cost(tokens_in, tokens_out), t)
            except OverBudget:
                allowed = False
        if allowed:
            served[user] += 1
            spent[user] += cost(tokens_in, tokens_out)
        else:
            refused[user] += 1
    return Day(rule, tuple(served), tuple(refused), tuple(spent))