"""Guards for prefix_cache.py: a kept start is reused whole or not at all, the bill adds up by hand, the rules behave as the lesson says."""
import os
from decimal import Decimal

import pytest

import prefix_cache as pc
from prefix_cache import RULES, Cache, bill_request, chat_bill, chat_cost, dollars, reuse, shared_start, tokens, view

A, B, C, D = RULES
EXAMPLE = chat_bill.EXAMPLE         # lesson 7: system 500, message 60, reply 200, 10 turns, $2 in and $10 out per million


def test_a_kept_start_is_reused_whole_or_not_at_all_and_what_follows_a_difference_does_not_count():
    """Guards the lesson's central claim with real text: a changed first line loses the whole start, the same line added at the end loses nothing
    before it, and a request that differs anywhere inside a kept start reads none of it (no partial reuse up to the first difference)."""
    earlier, new_question, clock_top, clock_end = (tokens(t, 0) for t in (pc.EARLIER, pc.PRESETS[0]["text"], pc.PRESETS[1]["text"], pc.PRESETS[2]["text"]))
    assert shared_start(earlier, clock_top) == 0
    assert shared_start(earlier, clock_end) == len(earlier)
    assert 0 < shared_start(earlier, new_question) < len(earlier) and earlier[-3:] != new_question[-3:]
    assert shared_start(new_question, earlier) == shared_start(earlier, new_question)
    assert reuse(earlier, clock_end) == len(earlier) and reuse(earlier, clock_top) == 0
    assert reuse(earlier, new_question) == 0                               # 41 tokens match, and none of them is reused
    assert (reuse([1, 2, 3], [1, 2, 3, 4]), reuse([1, 2, 3], [1, 2, 9]), reuse([1, 2, 3], [1, 2]), reuse([], [1])) == (3, 0, 0, 0)
    # the tokens after the first difference are equal here, and still add nothing
    one, two = [1, 2, 3, 4, 5, 6], [1, 2, 9, 4, 5, 6]
    assert shared_start(one, two) == 2
    for a, b in ((one, two), (earlier, new_question), (earlier, clock_end), ([], one), (one, one[:4])):      # an independent count, character by character
        assert shared_start(a, b) == len(os.path.commonprefix([a, b]))


def test_the_document_marker_counts_as_the_same_stand_in_tokens_wherever_it_matches():
    for doc in pc.DOC_LENGTHS:
        a, b = tokens(pc.EARLIER, doc), tokens(pc.PRESETS[0]["text"], doc)
        assert len(a) - len(tokens(pc.EARLIER, 0)) == doc
        assert shared_start(a, b) - shared_start(tokens(pc.EARLIER, 0), tokens(pc.PRESETS[0]["text"], 0)) == doc
    for doc in pc.DOC_LENGTHS:                  # the mark falls on a token boundary: every request that repeats the fixed part shares at least that many tokens
        a, mark = tokens(pc.EARLIER, doc), pc.mark_tokens(pc.EARLIER, doc)
        assert pc.shared_start(a, tokens(pc.PRESETS[0]["text"], doc)) >= mark and pc.shared_start(a, tokens(pc.PRESETS[2]["text"], doc)) >= mark
    assert pc.mark_at("x\n[document]\nlast") == 13 and pc.mark_at("a\nb\nlast") == 4 and pc.mark_at("one line") == 0 and pc.mark_at("[document] x") == 12
    ids = tokens(pc.EARLIER, 3)
    assert [s[0] for s in pc.segments(ids, 0, pc.merges()) if s[2]] == ["document, 3 tokens"]      # three stand-ins are drawn as one block
    assert [s[1] for s in pc.segments(ids, len(ids) - 1, pc.merges())].count("rest") == 1


def test_bills_by_hand():
    """1,200 tokens on rule A ($2 normal, $2.50 write, $0.10 read per million): a first request writes, a repeat reads."""
    assert bill_request(A, 1200, 0).units == 1200 * 2500 and dollars(bill_request(A, 1200, 0).units) == Decimal("0.003")
    assert bill_request(A, 1200, 1200).units == 1200 * 100 and dollars(bill_request(A, 1200, 1200).units) == Decimal("0.00012")
    assert bill_request(A, 1200, 1100).units == 1100 * 100 + 100 * 2500            # the 100 new tokens are written
    assert bill_request(A, 1200, 0).plain_units == 1200 * 2000
    # below the minimum nothing is stored or read, and the price is the plain price
    assert bill_request(A, 511, 511) == bill_request(A, 511, 0) and bill_request(A, 511, 0).units == 511 * 2000
    assert bill_request(D, 1000, 0).units == 1000 * 2000 and bill_request(D, 1024, 0).units == 1024 * 2500
    # a kept start under the minimum is not read, and the part to store is written instead
    assert bill_request(A, 1300, 200).read == 0 and bill_request(A, 1300, 200).written == 1300
    # a caller who marks only the fixed start pays the write premium on that start only
    marked = bill_request(A, 1200, 0, mark=1000)
    assert (marked.read, marked.written, marked.plain) == (0, 1000, 200)
    later = bill_request(A, 1200, 800, mark=1000)                                  # 800 read from a kept start, 200 more written up to the mark
    assert (later.read, later.written, later.plain) == (800, 200, 200)


def test_a_copy_lives_for_its_lifetime_after_its_last_use_and_a_read_renews_it():
    """Guards the five-minute rule: a request at minute 5 hits, minute 6 misses; requests every 4 minutes keep a copy alive forever."""
    prompt = [("t", i) for i in range(600)]
    for rule in (A, B, C, D):
        life = rule.lifetime
        for gap, hit in ((life, True), (life + 1, False)):
            cache = Cache(rule)
            cache.request(prompt, 0) if rule.minimum <= 600 else cache.request(prompt + prompt, 0)
            second = cache.request(prompt if rule.minimum <= 600 else prompt + prompt, gap)
            assert (second.read > 0) == hit, (rule.key, gap)
    cache = Cache(A)
    reads = [cache.request(prompt, minute).read for minute in (0, 4, 8, 12, 16)]
    assert reads == [0, 600, 600, 600, 600]
    cache = Cache(A)
    assert [cache.request(prompt, minute).read for minute in (0, 6, 12)] == [0, 0, 0]
    for rule in RULES:                                                              # the widget's view agrees: a stored start of exactly its lifetime is still read
        assert view(pc.EARLIER, pc.PRESETS[0]["text"], rule, 1200, rule.lifetime, True)["match"].startswith("1,")
        assert view(pc.EARLIER, pc.PRESETS[0]["text"], rule, 1200, rule.lifetime + 0.5, True)["match"] == "No token is reused."


def test_the_cache_reads_the_longest_kept_start_the_request_begins_with_and_nothing_from_one_it_only_partly_shares():
    cache = Cache(A)
    base, more = [("b", i) for i in range(700)], [("m", i) for i in range(300)]
    assert cache.request(base, 0).written == 700
    assert cache.request(base + more, 1).read == 700                                # two starts are kept now: 700 and 1,000 tokens
    longer = cache.request(base + more + [("x", i) for i in range(100)], 2)
    assert (longer.read, longer.written, longer.plain) == (1000, 100, 0)
    shorter = cache.request(base + [("y", i) for i in range(300)], 3)               # the 1,000-token start differs at token 700: the 700-token start still matches whole
    assert (shorter.read, shorter.written, shorter.plain) == (700, 300, 0)
    partly = cache.request(base[:650] + [("z", i) for i in range(200)], 4)           # 650 of 700 shared: all or nothing
    assert (partly.read, partly.written, partly.plain) == (0, 850, 0)
    marked = Cache(A)                                                                # a mark keeps only the start up to it, not the whole request
    assert marked.request(base + more, 0, mark=700).written == 700
    after = marked.request(base + [("y", i) for i in range(300)], 1)
    assert (after.read, after.written, after.plain) == (700, 300, 0)


def test_break_even_is_the_fewest_reads_that_repay_the_write():
    """Guards the break-even numbers (1, 2, 1, 1) against the simulation itself: n requests of one 1,200-token start, cached against not cached."""
    assert [pc.break_even_reads(r) for r in RULES] == [1, 2, 1, 1]
    for rule in RULES:
        prompt = [("doc", i) for i in range(1200)]
        def total(n):
            cache = Cache(rule)
            return sum(cache.request(prompt, k).units for k in range(n))
        extra = pc.break_even_reads(rule)
        assert total(1 + extra) <= (1 + extra) * 1200 * rule.units("price")          # one write and `extra` reads: no worse than not caching
        assert total(extra) > extra * 1200 * rule.units("price")                    # one read fewer: still behind
        assert total(1) == 1200 * rule.units("write") > 1200 * rule.units("price")   # a start used once costs the write premium


def test_example_chat_uncached_matches_lesson_7_and_cached_matches_the_closed_form():
    """Guards the chat bill: with caching off the sim must give lesson 7's $0.0546; with it on, turn 1 writes everything and every later turn
    reads the whole previous prompt and writes only the reply and message that are new."""
    off = chat_cost(EXAMPLE, A, cached=False)
    assert dollars(off["total"]) == chat_bill.bill(EXAMPLE, *chat_bill.EXAMPLE_PRICES).total == Decimal("0.0546")
    on = chat_cost(EXAMPLE, A)
    prompts = [560 + 260 * k for k in range(10)]
    assert (on["read"], on["written"], on["plain"]) == (sum(prompts[:-1]), sum(prompts) - sum(prompts[:-1]), 0) == (14400, 2900, 0)
    assert dollars(on["input"]) == Decimal("0.00869") and dollars(on["total"]) == Decimal("0.02869") and dollars(on["output"]) == Decimal("0.02")
    assert pc.percent_less(off["total"], on["total"]) == 47 and pc.percent_less(off["input"], on["input"]) == 75


def test_where_caching_saves_nothing_or_costs_more():
    """Guards the three cases the lesson names: a first line that changes every turn, turns spaced wider than the lifetime, and a start under the minimum."""
    off = chat_cost(EXAMPLE, A, cached=False)["total"]
    assert chat_cost(EXAMPLE, A, gap=1)["total"] < off
    assert chat_cost(EXAMPLE, A, gap=6)["read"] == 0 and chat_cost(EXAMPLE, A, gap=6)["total"] > off
    assert chat_cost(EXAMPLE, B, gap=6)["total"] < off                             # the 1-hour copy is still there
    stamped = chat_cost(EXAMPLE, A, stamp=True)
    assert stamped["read"] == 0 and stamped["total"] > chat_cost(EXAMPLE, A, cached=False, stamp=True)["total"]
    assert pc.percent_less(off, chat_cost(EXAMPLE, A, gap=6)["total"]) == -16
    short = chat_bill.Chat(system=100, message=20, reply=50, turns=5)               # no prompt reaches 512 tokens
    assert chat_cost(short, A)["total"] == chat_cost(short, A, cached=False)["total"]
    d = chat_cost(EXAMPLE, D)                                                      # rule D: the first two prompts are under 1,024 tokens
    assert d["plain"] == 560 + 820 and d["read"] == sum(560 + 260 * k for k in range(2, 9))


@pytest.mark.parametrize("key, every_minute, every_40_minutes, clock_first, uncached", [
    ("A", "0.02869", "0.06325", "0.063275", "0.0546"),
    ("B", "0.03304", "0.03304", "0.08924", "0.0546"),
    ("C", "0.0015065", "0.0031625", "0.00316375", "0.00273"),
    ("D", "0.031312", "0.06256", "0.062584", "0.0546"),
])
def test_chat_totals_of_the_four_rules_are_pinned(key, every_minute, every_40_minutes, clock_first, uncached):
    """Guards the output price and the lifetime of each rule: the lesson's table prints these totals, and a wrong output price (rule B 20 for 10,
    rule C 10 for 0.50, rule D 5 for 10) moved none of the other tests."""
    rule = pc.RULE_BY_KEY[key]
    assert dollars(chat_cost(EXAMPLE, rule, gap=1)["total"]) == Decimal(every_minute)
    assert dollars(chat_cost(EXAMPLE, rule, gap=40)["total"]) == Decimal(every_40_minutes)
    assert dollars(chat_cost(EXAMPLE, rule, gap=1, stamp=True)["total"]) == Decimal(clock_first)
    assert dollars(chat_cost(EXAMPLE, rule, cached=False)["total"]) == Decimal(uncached)


def test_a_reused_start_skips_its_reading_time():
    """Guards the best-case figures (lesson 23's model): reading 8,020 tokens takes 412.8 ms at the chip's peak; with 8,000 reused only 20 are read."""
    for total, read, before, after in ((8020, 8000, "412.8", "7.9"), (2900, 2640, "149.3", "13.4"), (560, 0, "28.8", "28.8")):
        full, rest = pc.first_word_seconds(total, read)
        assert (f"{full * 1000:.1f}", f"{rest * 1000:.1f}") == (before, after)
    assert pc.first_word_seconds(100, 100)[1] == pc.prefill_sim.prefill(1).seconds        # the newest token is always read


NEW_QUESTION, CLOCK_TOP, CLOCK_END = (p["text"] for p in pc.PRESETS)


@pytest.mark.parametrize("text, key, doc, age, marked, expect", [
    (NEW_QUESTION, "A", 1200, 1, True, ["1,236 of 1,254 tokens are reused from the stored start", "1,236 tokens at $0.10 (read), 0 at $2.50 (write), 18 at $2.00 (normal)",
                                        "$0.000160 with caching, $0.002508 without: 94% less", "begins with the whole stored start"]),
    (CLOCK_TOP, "A", 1200, 1, True, ["No token is reused.", "0 tokens at $0.10 (read), 1,255 at $2.50 (write), 17 at $2.00 (normal)",
                                     "$0.003172 with caching, $0.002544 without: 25% more", "very first token differs"]),
    (CLOCK_END, "A", 1200, 1, True, ["1,236 of 1,272 tokens are reused", "36 at $2.00 (normal)", "92% less", "The other 36 tokens are read again"]),
    (NEW_QUESTION, "A", 1200, 1, False, ["No token is reused: without a mark, only the whole earlier request is stored.",
                                         "0 tokens at $0.10 (read), 1,254 at $2.50 (write), 0 at $2.00 (normal)", "$0.003135 with caching, $0.002508 without: 25% more",
                                         "Nothing is reused, and the whole prompt is written as a new copy. Without a mark, only the whole earlier request is stored, and the new "
                                         "request stops matching the stored start after 1,241 of its 1,254 tokens."]),
    (pc.EARLIER, "A", 1200, 1, False, ["1,253 of 1,253 tokens are reused", "Nothing is left to read again."]),
    (pc.EARLIER.rsplit(" ", 1)[0], "A", 1200, 1, False, ["and the new request is shorter than the stored start."]),
    (pc.EARLIER.split("\n[document]")[0] + "\n[document]", "A", 1200, 1, True, ["The new request is shorter than the stored start. A stored start is reused whole"]),
    (CLOCK_END, "A", 1200, 1, False, ["1,253 of 1,272 tokens are reused", "93% less"]),
    (pc.EARLIER.replace("two short", "three short"), "A", 1200, 1, True, ["No token is reused.", "stops matching the stored start after 24 of its 1,252 tokens",
                                                                           "reused whole or not at all", "1,235 at $2.50 (write)"]),
    (NEW_QUESTION, "A", 1200, 6, True, ["25% more", "older than 5 minutes", "stored start is gone"]),
    (NEW_QUESTION, "B", 1200, 6, True, ["Billed: 1,236 tokens at $0.10 (read), 0 at $4.00 (write), 18 at $2.00 (normal)"]),
    (NEW_QUESTION, "A", 0, 1, True, ["shorter than the 512-token minimum, so nothing is stored or read", "about the same"]),
    (NEW_QUESTION, "D", 400, 1, True, ["shorter than the 1,024-token minimum"]),
    ("[document]\n[document]", "A", 400, 1, True, ["The earlier request stores 436 tokens, under the 512-token minimum"]),
])
def test_widget_words_say_what_the_settings_do(text, key, doc, age, marked, expect):
    """Guards the sentences a reader sees: each case moves one setting or the mark and checks the reused count, the tokens at each price, the cost and the reason."""
    v = view(pc.EARLIER, text, pc.RULE_BY_KEY[key], doc, age, marked)
    said = " ".join((v["match"], v["billed"], v["cost"], v["why"]))
    for part in expect:
        assert part in said, (part, said)
    assert view(pc.EARLIER, "", A, 1200, 1, True)["match"] == "Type a new request to compare."