"""Guards for bill.py: the closed form must equal the turn-by-turn simulation, and trimming and
the reasoning and retry terms must add exactly what the lesson says."""
import itertools
from decimal import Decimal
from fractions import Fraction

import pytest

from bill import Chat, bill, closed_form, naive_bill, simulate, turn_cost, turn_parts

EXAMPLE = Chat(system=500, message=60, reply=200, turns=10)


def grid():
    for s, m, r, n, think, keep in itertools.product(
            (0, 500), (0, 60), (0, 200), range(1, 13), (0, 300), (None, 0, 1, 3, 20)):
        yield Chat(system=s, message=m, reply=r, turns=n, reasoning=think, keep=keep)


def test_closed_form_equals_simulating_every_turn():
    """Guards the algebra: a slip in n(n-1)/2 or in the trimmed sum shows up on some small chat."""
    count = 0
    for chat in grid():
        assert closed_form(chat) == simulate(chat), chat
        count += 1
    assert count == 2 * 2 * 2 * 12 * 2 * 5


def test_the_per_turn_parts_add_up_to_the_same_totals():
    for chat in grid():
        parts = [turn_parts(chat, k) for k in range(1, chat.turns + 1)]
        assert (sum(p.input for p in parts), sum(p.output for p in parts)) == closed_form(chat), chat


def test_example_chat_by_hand():
    """10 turns: input 500+60 at turn 1, then 260 more each turn; output 200 each turn."""
    assert [turn_parts(EXAMPLE, k).input for k in (1, 2, 3, 10)] == [560, 820, 1080, 2900]
    assert closed_form(EXAMPLE) == (5600 + 260 * 45, 2000) == (17300, 2000)
    b = bill(EXAMPLE, 2, 10)
    assert (b.input_cost, b.output_cost, b.total) == (Decimal("0.0346"), Decimal("0.02"), Decimal("0.0546"))


def test_untrimmed_input_grows_with_the_square_of_the_turns():
    """The first difference rises by m+r each turn, so the second difference is constant: a parabola."""
    def total(n):
        return closed_form(Chat(500, 60, 200, n))[0]
    firsts = [total(n) - total(n - 1) for n in range(2, 30)]
    assert firsts[0] == 500 + 60 + 260
    assert {b - a for a, b in zip(firsts, firsts[1:])} == {260}


def test_trimmed_cost_grows_linearly_once_the_window_is_full():
    """Guards the claim 'trim to N turns and the bill stops growing faster': every extra turn adds
    the same S + m + N(m+r) input tokens."""
    for keep in (0, 1, 3, 8):
        def total(n):
            return closed_form(Chat(500, 60, 200, n, keep=keep))[0]
        steps = {total(n) - total(n - 1) for n in range(keep + 2, 60)}
        assert steps == {500 + 60 + keep * 260}, keep


def test_trimming_changes_nothing_until_the_window_fills_and_never_costs_more():
    for n in range(1, 10):        # a window of 8 earlier turns holds every earlier turn of a chat up to 9 turns
        assert closed_form(Chat(500, 60, 200, n, keep=8)) == closed_form(Chat(500, 60, 200, n))
    for chat in grid():
        untrimmed = Chat(chat.system, chat.message, chat.reply, chat.turns, chat.reasoning)
        assert closed_form(chat)[0] <= closed_form(untrimmed)[0]
        assert closed_form(chat)[1] == closed_form(untrimmed)[1]


def test_reasoning_tokens_add_output_cost_only():
    """Guards the lesson sentence: hidden reasoning is billed as output, once per reply, and is
    not added to the input here (the model assumes it is not re-sent)."""
    plain, thinking = Chat(500, 60, 200, 10), Chat(500, 60, 200, 10, reasoning=800)
    a, b = bill(plain, 2, 10), bill(thinking, 2, 10)
    assert b.input_tokens == a.input_tokens
    assert b.output_tokens == a.output_tokens + 10 * 800
    assert b.total - a.total == Decimal(10 * 800) * 10 / 1_000_000 == Decimal("0.08")


def test_retries_add_their_share_of_every_call_again():
    base = bill(EXAMPLE, 2, 10)
    assert base.retry_cost == 0
    assert bill(EXAMPLE, 2, 10, retry_rate="0.1").retry_cost == base.total / 10
    assert bill(EXAMPLE, 2, 10, retry_rate="0.1").total == base.total * Decimal("1.1")


def test_retry_rate_is_the_expected_extra_cost_over_every_choice_of_retried_turns():
    """Guards 'the cost of retries is the retry rate times the bill': add up every way of choosing
    which turns of a 6-turn chat are made twice, each way weighted by its probability, exactly."""
    chat, q = Chat(500, 60, 200, 6, reasoning=100), Fraction(1, 4)
    per_turn = [Fraction(turn_cost(chat, k, 2, 10)) for k in range(1, 7)]
    expected = Fraction(0)
    for picked in itertools.product((False, True), repeat=6):
        p = Fraction(1)
        for flag in picked:
            p *= q if flag else 1 - q
        expected += p * sum(c * (2 if flag else 1) for c, flag in zip(per_turn, picked))
    assert expected == sum(per_turn) * (1 + q)
    assert Decimal(expected.numerator) / Decimal(expected.denominator) == bill(chat, 2, 10, retry_rate="0.25").total


def test_the_naive_estimate_misses_exactly_the_system_prompt_and_the_history():
    """Guards the lesson's opening claim: ignoring what rides along makes the estimate 2.6 times too low."""
    assert naive_bill(EXAMPLE, 2, 10) == Decimal("0.0212")
    long = Chat(system=500, message=60, reply=200, turns=30)
    assert naive_bill(long, 2, 10) == Decimal("0.0636")
    # hidden reasoning is output, so the naive estimate counts it: 10 turns x 600 more tokens at $10 per million
    assert naive_bill(Chat(system=500, message=60, reply=200, turns=10, reasoning=600), 2, 10) == Decimal("0.0812")
    for chat in (EXAMPLE, long, Chat(system=0, message=60, reply=200, turns=1)):
        tokens_in, _ = closed_form(chat)
        missed = Decimal(tokens_in - chat.turns * chat.message) * 2 / 1_000_000
        assert bill(chat, 2, 10).total - naive_bill(chat, 2, 10) == missed
    # a one-turn chat with no system prompt has nothing riding along, so the estimate is exact
    assert naive_bill(Chat(system=0, message=60, reply=200, turns=1), 2, 10) == bill(Chat(system=0, message=60, reply=200, turns=1), 2, 10).total


def test_the_check_yourself_answers():
    """Guards the numbers printed in the lesson's questions, so an edit to either side shows up here."""
    small = Chat(system=100, message=10, reply=20, turns=4)
    assert closed_form(small) == (4 * 110 + 30 * 6, 80) == (620, 80)
    assert closed_form(Chat(system=100, message=10, reply=20, turns=4, keep=1)) == (4 * 110 + 30 * 3, 80) == (530, 80)
    assert bill(small, 2, 10).total == Decimal("0.00204")
    thinker = Chat(system=0, message=0, reply=150, turns=1, reasoning=600)
    assert bill(thinker, 2, 10).output_cost == Decimal("0.0075") == 5 * bill(Chat(system=0, message=0, reply=150, turns=1), 2, 10).output_cost
    eight = Chat(system=200, message=30, reply=70, turns=8)
    assert closed_form(eight)[0] == 4640 and closed_form(Chat(system=200, message=30, reply=70, turns=8, keep=2))[0] == 3140


@pytest.mark.parametrize("kwargs", [
    dict(system=-1, message=1, reply=1, turns=1), dict(system=0, message=-1, reply=1, turns=1),
    dict(system=0, message=1, reply=-1, turns=1), dict(system=0, message=1, reply=1, turns=0),
    dict(system=0, message=1, reply=1, turns=1, reasoning=-1), dict(system=0, message=1, reply=1, turns=1, keep=-1),
])
def test_nonsense_sizes_are_refused(kwargs):
    with pytest.raises(ValueError):
        Chat(**kwargs)


def test_negative_prices_and_retry_rates_are_refused():
    for bad in (dict(price_in=-1, price_out=1), dict(price_in=1, price_out=-1), dict(price_in=1, price_out=1, retry_rate=-1)):
        with pytest.raises(ValueError):
            bill(EXAMPLE, **bad)