"""Tests for budget.py and inflight.py. Each test names the failure it guards."""
import re
import textwrap
import threading
import time
from pathlib import Path

import pytest

from budget import (DAY, MAX_REPLY, Budget, CheckOnlyBudget, OverBudget, cost, dollars, handle_chat, how_long, input_tokens_upper_bound, refusal)
from inflight import BURST_INPUT, BURST_OUTPUT, CALLS, fire
from sim import simulate_day

LIMIT = 250_000  # $0.25
REQUEST = {"messages": [{"role": "user", "content": "hello"}]}


def ok_upstream(tokens_in=100, tokens_out=50):
    return lambda request, max_reply: {"text": "hi", "usage": {"input_tokens": tokens_in, "output_tokens": tokens_out}}


def test_one_call_can_cost_hundreds_of_times_another():
    """The lesson's reason a request limit is not a cost limit: a short question and a long document, same price list."""
    short, long_ = cost(60, 30), cost(50_000, 2_000)
    assert (short, long_) == (420, 120_000)
    assert 285 < long_ / short < 286


def test_reserve_allows_an_exact_fit_and_refuses_one_micro_dollar_more():
    """Guards >= instead of > in the check: a user could never spend the last of the budget."""
    budget = Budget(1000)
    budget.reserve("a", 1000, 0)
    with pytest.raises(OverBudget):
        budget.reserve("a", 1, 0)
    with pytest.raises(OverBudget):
        Budget(1000).reserve("a", 1001, 0)


def test_settle_charges_the_real_cost_and_gives_back_the_rest():
    budget = Budget(LIMIT)
    hold = budget.reserve("a", 50_000, 0)
    assert budget.used("a", 0) == 50_000
    budget.settle(hold, 48_000, 0)
    assert budget.used("a", 0) == 48_000


def test_a_reply_that_costs_more_than_the_set_aside_is_charged_in_full():
    """Guards a settle that clamps to the hold: a wrong estimate would then be free, and the budget would drift low."""
    budget = Budget(100)
    hold = budget.reserve("a", 60, 0)
    budget.settle(hold, 90, 0)
    assert budget.used("a", 0) == 90
    budget.settle(budget.reserve("a", 10, 0), 30, 0)  # 90 + 10 fits; the real cost takes the user to 120
    assert budget.used("a", 0) == 120
    with pytest.raises(OverBudget):
        budget.reserve("a", 1, 0)


def test_release_gives_back_the_whole_set_aside():
    """Guards a failed call that keeps its hold: the user would pay for nothing."""
    budget = Budget(LIMIT)
    budget.release(budget.reserve("a", 50_000, 0), 0)
    assert budget.used("a", 0) == 0


def test_users_have_separate_budgets():
    """Guards one shared counter: one heavy user would lock everyone out."""
    budget = Budget(100)
    budget.reserve("heavy", 100, 0)
    budget.reserve("light", 100, 0)
    with pytest.raises(OverBudget):
        budget.reserve("heavy", 1, 0)


def test_daily_budget_restarts_at_midnight_and_a_hold_from_yesterday_is_ignored():
    """Guards a settle that lands after the reset: it would subtract yesterday's hold from today's fresh count and go negative."""
    budget = Budget(100)
    hold = budget.reserve("a", 50, DAY - 5)
    budget.settle(hold, 80, DAY + 3)  # a call that started before midnight ends after it
    assert budget.used("a", DAY + 3) == 0
    old = budget.reserve("b", 50, DAY - 5)
    budget.reserve("b", 40, DAY + 1)
    budget.release(old, DAY + 3)  # and a call that failed after midnight gives nothing back from today's count
    assert budget.used("b", DAY + 3) == 40
    today = budget.reserve("c", 50, DAY + 10)  # a call set aside and settled on the same day after the first midnight still settles
    budget.settle(today, 30, DAY + 20)
    assert budget.used("c", DAY + 20) == 30
    budget.settle(today, 0, DAY + 30)  # a hold settled twice must not take the count below zero and bank free credit
    assert budget.used("c", DAY + 30) == 0


def test_steady_refill_drains_by_limit_over_a_day_and_says_how_long_to_wait():
    budget = Budget(86_400, refill="steady")  # 1 micro-dollar back per second
    budget.reserve("a", 86_400, 0)
    assert budget.used("a", 1000) == 85_400
    with pytest.raises(OverBudget) as refused:
        budget.reserve("a", 2000, 1000)  # 1,000 is free, so the call is 1,000 too big
    assert refused.value.wait == 1000  # at 1 micro-dollar back per second
    budget.reserve("a", 1000, 1000)
    assert budget.used("a", 1000 + DAY) == 0
    budget.reserve("a", 86_400, 1000 + 10 * DAY)  # ten idle days bank no credit: the whole budget is free and not a unit more
    with pytest.raises(OverBudget):
        budget.reserve("a", 1, 1000 + 10 * DAY)


def test_a_call_bigger_than_the_whole_budget_is_told_to_send_less_not_to_wait():
    """Guards a message that says "try again later" for a call that can never fit."""
    status, body = handle_chat(Budget(LIMIT), ok_upstream(), "a", {"messages": [{"role": "user", "content": "x" * 200_000}]}, 0)
    assert status == 429 and body["error"] == "request_too_big" and "resets_at" not in body
    assert "Send less text" in body["message"]


def test_refused_call_gets_a_clear_message_and_the_refill_time_not_a_bare_429():
    """Lesson 10: a used-up quota is not a retry case. The body says so in words and in a code the app can test; there is no retry hint."""
    budget = Budget(LIMIT)
    now = 14 * 3600 + 40 * 60  # 14:40, so 9 h 20 min to midnight
    budget.reserve("a", 245_000, now)
    status, body = handle_chat(budget, ok_upstream(), "a", REQUEST, now)
    assert status == 429
    assert body["error"] == "budget_used_up"
    assert body["resets_at"] == DAY
    assert body["message"] == "Your AI budget has $0.01 left and this call needs up to $0.01. It refills in 9 h 20 min."
    assert not any("retry" in key for key in body)


def test_a_user_who_is_not_signed_in_gets_401_and_costs_nothing():
    calls = []
    status, body = handle_chat(Budget(LIMIT), lambda *a: calls.append(a), None, REQUEST, 0)
    assert (status, body["error"], calls) == (401, "sign_in", [])


def test_an_answered_call_is_charged_the_usage_numbers_not_the_estimate():
    budget = Budget(LIMIT)
    status, body = handle_chat(budget, ok_upstream(100, 50), "a", REQUEST, 0)
    assert (status, body["charged"]) == (200, cost(100, 50))
    assert budget.used("a", 0) == cost(100, 50)


def test_a_failed_upstream_call_costs_the_user_nothing_and_a_timed_out_one_keeps_its_set_aside():
    """Guards a missing release on the error path (each outage would eat the user's budget) and a release on a timeout (the provider may still bill that call)."""
    budget = Budget(LIMIT)

    def broken(request, max_reply):
        raise ConnectionError("down")

    def slow(request, max_reply):
        raise TimeoutError("no answer")

    status, body = handle_chat(budget, broken, "a", REQUEST, 0)
    assert status == 502 and body["error"] == "upstream_failed" and "Nothing was charged" not in body["message"]
    assert budget.used("a", 0) == 0
    status, body = handle_chat(budget, slow, "a", REQUEST, 0)
    assert status == 502 and body["error"] == "upstream_timeout" and "Nothing was charged" not in body["message"]
    assert budget.used("a", 0) == cost(input_tokens_upper_bound(REQUEST["messages"]), MAX_REPLY)


def test_the_set_aside_is_input_plus_the_reply_cap_priced_as_output():
    seen = {}
    budget = Budget(LIMIT)
    handle_chat(budget, lambda request, max_reply: seen.setdefault("held", budget.used("a", 0)) and ok_upstream()(request, max_reply),
                "a", REQUEST, 0, count_tokens=lambda messages: 700)
    assert seen["held"] == cost(700, MAX_REPLY)


def test_the_upper_bound_never_falls_below_the_real_count():
    """A token covers at least one byte of text, so bytes plus the per-message extra is a safe ceiling (English runs near 4 bytes a token)."""
    messages = [{"role": "user", "content": "naïve café 你好"}, {"role": "assistant", "content": "ok"}]
    bound = input_tokens_upper_bound(messages)
    assert bound == len("naïve café 你好".encode()) + len(b"ok") + 2 * 10
    assert bound >= len("naïve café 你好")


@pytest.mark.parametrize("seconds, text", [(9 * 3600 + 20 * 60, "9 h 20 min"), (1, "1 min"), (3601, "1 h 1 min"), (59 * 60, "59 min"), (3600, "1 h 0 min")])
def test_how_long_rounds_up_to_a_whole_minute(seconds, text):
    assert how_long(seconds) == text


@pytest.mark.parametrize("micro, text", [(250_000, "$0.25"), (4_999, "$0.00"), (5_000, "$0.01"), (9_925_000, "$9.93"), (1_000_000, "$1.00")])
def test_dollars_round_halves_up_to_whole_cents(micro, text):
    assert dollars(micro) == text


def test_twenty_calls_at_once_all_pass_a_check_made_before_any_is_charged():
    """The in-flight overrun. Without a set-aside every call sees an empty budget; with it only five fit."""
    each = cost(BURST_INPUT, BURST_OUTPUT)
    assert each == 48_000
    statuses, charged = fire(CheckOnlyBudget(LIMIT))
    assert (statuses.count(200), statuses.count(429), charged) == (CALLS, 0, CALLS * each)  # 20 calls, $0.96, 3.84 times the budget
    statuses, charged = fire(Budget(LIMIT))
    assert (statuses.count(200), statuses.count(429), charged) == (5, 15, 5 * each)  # 5 fit: 5 x $0.05 set aside = $0.25
    assert charged <= LIMIT


def test_two_server_processes_with_their_own_counts_each_allow_the_whole_budget():
    """Why the state must live in one shared store once there is more than one process (two Budget objects stand in for two processes)."""
    first, second = fire(Budget(LIMIT)), fire(Budget(LIMIT))
    assert first[0].count(200) + second[0].count(200) == 10
    assert first[1] + second[1] == 480_000 > 1.9 * LIMIT


def test_a_low_estimate_lets_more_calls_in_and_the_budget_is_overshot():
    """Guards the claim in the lesson: the set-aside is only as good as the estimate. Half the real input size admits 8 calls, not 5."""
    statuses, charged = fire(Budget(LIMIT), estimate=BURST_INPUT // 2)
    assert (statuses.count(200), charged) == (8, 384_000)
    assert charged > 1.5 * LIMIT


def test_the_check_and_the_set_aside_are_one_step_under_threads():
    """Guards a missing lock: a slow read widens the gap between the check and the set-aside, so threads that skip the lock all see an empty budget."""

    class SlowRead(Budget):
        def _fresh(self, user, now):
            fresh = super()._fresh(user, now)
            time.sleep(0.01)  # a real store takes time to answer
            return fresh

    budget, wins = SlowRead(1), []

    def grab():
        try:
            budget.reserve("a", 1, 0)
            wins.append(1)
        except OverBudget:
            pass

    threads = [threading.Thread(target=grab) for _ in range(8)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()
    assert len(wins) == 1


def test_refusal_message_names_what_is_left_what_is_needed_and_the_refill():
    refused = OverBudget(cost(50_000, MAX_REPLY), 30_000, 3600)  # a 50,000-token input with the reply cap
    assert refusal(refused, LIMIT, 0) == {"error": "budget_used_up", "resets_at": 3600,
                                         "message": "Your AI budget has $0.03 left and this call needs up to $0.11. It refills in 1 h 0 min."}
    over = Budget(LIMIT)
    over.settle(over.reserve("a", 1, 0), 300_000, 0)  # a reply that overran its set-aside: more than the whole budget is used
    with pytest.raises(OverBudget) as refused:
        over.reserve("a", 1, 0)
    assert "$0.00 left" in refusal(refused.value, LIMIT, 0)["message"]  # never a negative amount such as "$-1.95 left"


HERE = Path(__file__).parent
LESSON = (HERE.parent / "lesson.md").read_text()


def test_every_python_block_in_the_lesson_is_a_verbatim_excerpt_of_the_source():
    """The lesson quotes these files. If the code changes and the lesson does not, the reader sees code that was never tested."""
    sources = [(HERE / name).read_text().splitlines() for name in ("budget.py", "inflight.py", "sim.py")]
    blocks = re.findall(r"^```python\n(.*?)^```", LESSON, flags=re.S | re.M)
    assert blocks, "the lesson has no python blocks to check"
    for block in blocks:
        lines = block.rstrip("\n").splitlines()
        windows = (textwrap.dedent("\n".join(source[i:i + len(lines)])) for source in sources for i in range(len(source) - len(lines) + 1))
        assert textwrap.dedent("\n".join(lines)) in windows, f"not in the source:\n{block}"


def test_the_lesson_quotes_the_numbers_the_code_gives():
    """Guards the prose drifting from the code: the excerpt test above covers the code blocks and not the numbers around them."""
    none, budget = simulate_day("none"), simulate_day("budget")
    heavy_calls, all_calls = none.served[0], sum(none.served)
    check_only, set_aside, half = fire(CheckOnlyBudget(LIMIT)), fire(Budget(LIMIT)), fire(Budget(LIMIT), estimate=BURST_INPUT // 2)
    first, second = fire(Budget(LIMIT)), fire(Budget(LIMIT))  # two server processes, each with its own count
    quotes = [
        f"costs ${cost(60, 30) / 1e6:.5f}", f"costs {dollars(cost(50_000, 2_000))}", f"about {round(cost(50_000, 2_000) / cost(60, 30))} times as much",
        f"that user is {none.percent(0)}% of a {dollars(none.bill)} bill", f"cuts the bill to {dollars(budget.bill)} and the heavy user's share to {budget.percent(0)}%",
        f"refuses none of the nine other users' {none.ordinary_asked} calls", f"makes {heavy_calls} of the day's {all_calls} calls, {round(100 * heavy_calls / all_calls)}%",
        f"is refused {budget.refused[0]} of {heavy_calls} calls", f"The bill is {dollars(check_only[1])}", f"the bill is {dollars(set_aside[1])}",
        f"the bill is ${half[1] / 1e6:.3f}", f"{set_aside[0].count(200)} calls each hold", f"{half[0].count(200)} calls fit",
        f"{first[0].count(200) + second[0].count(200)} get in and {dollars(first[1] + second[1])} is charged", f"turns away {simulate_day('budget', budget=50_000).ordinary_refused} of their calls",
        refusal(OverBudget(cost(50_000, MAX_REPLY), 30_000, 3600), LIMIT, 0)["message"],
        f"the nine ordinary users lose {simulate_day('rpm', rpm=1).ordinary_refused} calls",
    ]
    missing = [q for q in quotes if q not in LESSON]
    assert not missing, missing