"""What an AI chat bill is made of, turn by turn, and a closed form for the total.

Python standard library only. Run `python3 show_numbers.py` to print every number the
lesson quotes.

The model: every turn the app sends the provider the system prompt, the earlier turns
(unless trimmed) and the new message, and the provider bills those as input tokens. It
then bills the reply, plus any hidden reasoning tokens, as output tokens. The two kinds
of token have two prices. Sizes are in tokens; prices are dollars per million tokens.
Reasoning tokens are assumed not to be sent back as history. Some providers do send
them back (the lesson cites which), and then the history part would be larger.
"""
from dataclasses import dataclass
from decimal import Decimal
from typing import Optional

PER = Decimal(1_000_000)
DAYS_PER_MONTH = 30


@dataclass(frozen=True)
class Chat:
    system: int                  # system prompt, tool definitions and pasted documents: sent on every call
    message: int                 # tokens in each new user message
    reply: int                   # visible tokens in each reply
    turns: int                   # user messages in the chat
    reasoning: int = 0           # hidden reasoning tokens per reply, billed as output
    keep: Optional[int] = None   # send only the last `keep` earlier turns; None sends them all

    def __post_init__(self):
        if self.turns < 1 or min(self.system, self.message, self.reply, self.reasoning) < 0:
            raise ValueError("turns must be at least 1 and sizes cannot be negative")
        if self.keep is not None and self.keep < 0:
            raise ValueError("keep cannot be negative")


@dataclass(frozen=True)
class Turn:
    system: int
    history: int
    message: int
    reply: int
    reasoning: int

    @property
    def input(self):
        return self.system + self.history + self.message

    @property
    def output(self):
        return self.reply + self.reasoning


def turn_parts(chat, k):
    """The tokens of turn k (counting from 1), split the way the widget's bar splits them."""
    earlier = k - 1 if chat.keep is None else min(k - 1, chat.keep)
    return Turn(chat.system, earlier * (chat.message + chat.reply), chat.message, chat.reply, chat.reasoning)


def simulate(chat):
    """Input and output tokens of the whole chat, found by walking through every turn."""
    history = []  # earlier message sizes, oldest first
    total_in = total_out = 0
    for _ in range(chat.turns):
        kept = history
        if chat.keep is not None:
            kept = history[max(0, len(history) - 2 * chat.keep):]
        total_in += chat.system + sum(kept) + chat.message
        total_out += chat.reply + chat.reasoning
        history += [chat.message, chat.reply]
    return total_in, total_out


def closed_form(chat):
    """The same two totals without the loop.

    Untrimmed, turn k re-sends k-1 earlier turns, and 0 + 1 + ... + (n-1) = n(n-1)/2.
    Trimmed to N turns, the first N+1 turns re-send 0, 1, ..., N earlier turns and every
    later turn re-sends N.
    """
    n = chat.turns
    t = n - 1 if chat.keep is None else min(n - 1, chat.keep)
    earlier_turns = t * (t + 1) // 2
    if chat.keep is not None:
        earlier_turns += chat.keep * (n - 1 - t)
    fixed = n * (chat.system + chat.message)
    total_in = fixed + (chat.message + chat.reply) * earlier_turns
    return total_in, n * (chat.reply + chat.reasoning)


@dataclass(frozen=True)
class Bill:
    input_tokens: int
    output_tokens: int
    input_cost: Decimal
    output_cost: Decimal
    retry_cost: Decimal

    @property
    def total(self):
        return self.input_cost + self.output_cost + self.retry_cost


def _price(value, name):
    value = Decimal(str(value))
    if value < 0:
        raise ValueError(f"{name} cannot be negative")
    return value


def turn_cost(chat, k, price_in, price_out):
    """Dollars for turn k: input at the input price plus output at the output price."""
    parts = turn_parts(chat, k)
    return (parts.input * _price(price_in, "price_in") + parts.output * _price(price_out, "price_out")) / PER


def bill(chat, price_in, price_out, retry_rate=0):
    """Dollars for the whole chat. Prices are dollars per million tokens.

    retry_rate is the average number of extra full calls per call: 0.1 means one call in ten
    is made twice and both are billed. A retry sends the same input again, so it adds
    retry_rate times the whole bill.
    """
    price_in, price_out, retry_rate = _price(price_in, "price_in"), _price(price_out, "price_out"), _price(retry_rate, "retry_rate")
    tokens_in, tokens_out = closed_form(chat)
    cost_in = tokens_in * price_in / PER
    cost_out = tokens_out * price_out / PER
    retry_cost = (cost_in + cost_out) * retry_rate
    return Bill(tokens_in, tokens_out, cost_in, cost_out, retry_cost)


def naive_bill(chat, price_in, price_out):
    """The estimate most people make first: what you type, and what comes back, each billed once."""
    tokens_in = chat.turns * chat.message
    tokens_out = chat.turns * (chat.reply + chat.reasoning)
    cost_in = tokens_in * _price(price_in, "price_in")
    cost_out = tokens_out * _price(price_out, "price_out")
    return (cost_in + cost_out) / PER


# The chat the lesson and its widget start from. The prices are example figures, not a quote.
EXAMPLE = Chat(system=500, message=60, reply=200, turns=10)
EXAMPLE_PRICES = (2, 10)         # dollars per million input tokens, dollars per million output tokens
EXAMPLE_KEEP = 3                 # the "last N turns" the widget's trim box starts at
EXAMPLE_CHATS_PER_DAY = 100