import random

import pytest

from context_window import attention_pairs, count_tokens, total_tokens, tokens_read, truncate


def msg(role, text):
    return {"role": role, "content": text}


# The lesson's worked example. Word counts, worked out by hand: 3, 4, 2, 3, 2, 5.
CHAT = [
    msg("system", "Reply in French."),
    msg("user", "My cat is Biscuit."),
    msg("assistant", "Joli nom."),
    msg("user", "I like tea."),
    msg("assistant", "Moi aussi."),
    msg("user", "What is my cat called?"),
]


def test_counter_is_words():
    assert [count_tokens(m["content"]) for m in CHAT] == [3, 4, 2, 3, 2, 5]
    assert count_tokens("") == 0


def test_worked_example_pinned_system_drops_the_cat():
    """Oldest dropped first: the message with the cat's name goes, the pinned system message stays.
    The walk keeps messages 6 and 5, stops at message 4 (it does not fit), then drops the leading
    reply 5, whose question 4 was cut. Message 3 (2 tokens) would fit in the spare room, but keeping
    it would leave a hole."""
    kept = truncate(CHAT, 12)
    assert kept == [CHAT[0], CHAT[5]]
    assert total_tokens(kept) == 3 + 5 == 8


def test_worked_example_without_pin_drops_the_system_message_too():
    """The walk keeps messages 6, 5, 4 and 3 (12 tokens); reply 3 has lost its question 2, so it goes."""
    kept = truncate(CHAT, 12, keep_system=False)
    assert kept == CHAT[3:]
    assert total_tokens(kept) == 3 + 2 + 5 == 10


def test_a_reply_whose_question_was_cut_is_dropped_but_a_whole_chat_is_sent_whole():
    """Guards the old behaviour, where the window could start with a reply and no question."""
    chat = [msg("assistant", "Hello there"), msg("user", "Hi"), msg("assistant", "Hello again"), msg("user", "Bye")]
    assert truncate(chat, 6) == chat                    # nothing cut: the opening reply stays
    assert truncate(chat, 4) == chat[1:]                # cut, but the run starts with a question: nothing more to drop
    assert truncate(chat, 3) == [chat[3]]               # run = [reply, Bye]; the reply lost its question "Hi"
    two_replies = [msg("user", "q q q q q q"), msg("assistant", "a"), msg("assistant", "b"), msg("user", "c")]
    assert truncate(two_replies, 3) == [two_replies[3]]  # both leading replies go, not just one
    with pytest.raises(ValueError, match="newest"):      # only a reply fits, and its question was cut
        truncate([msg("user", "q q q q q q"), msg("assistant", "a")], 3)


def test_everything_fits_so_nothing_is_dropped():
    assert truncate(CHAT, 19) == CHAT
    assert truncate(CHAT, 1000) == CHAT


def test_no_holes_a_big_message_blocks_older_small_ones():
    """A greedy 'skip what does not fit' would keep A and leave a gap in the conversation."""
    chat = [msg("system", "S"), msg("user", "A"), msg("user", "b b b b b"), msg("user", "C"), msg("user", "D")]
    assert truncate(chat, 6) == [chat[0], chat[3], chat[4]]


def test_a_different_counter_changes_what_fits():
    """The limit is only meaningful under the counter that measured it."""
    chat = [msg("user", "aaaa"), msg("user", "bb"), msg("user", "cc")]
    assert truncate(chat, 5) == chat                     # words: 1 + 1 + 1 = 3 fits
    assert truncate(chat, 5, counter=len) == chat[1:]    # characters: 4 + 2 + 2 = 8 > 5, so the oldest goes
    assert truncate(chat, 8, counter=len) == chat


def test_per_message_overhead_counts_against_the_limit():
    chat = [msg("user", "a b"), msg("user", "c d"), msg("user", "e f")]
    assert truncate(chat, 9, per_message=1) == chat      # 3 + 3 + 3
    assert truncate(chat, 8, per_message=1) == chat[1:]  # one more word's worth would not fit


def test_errors_are_loud():
    with pytest.raises(ValueError, match="pinned"):
        truncate(CHAT, 2)
    with pytest.raises(ValueError, match="newest"):
        truncate(CHAT, 7)   # 3 pinned + 5 for the newest message = 8 > 7
    assert truncate([], 5) == []


def test_never_over_the_limit_and_greedy_from_the_newest():
    """Randomized: within the limit, a pinned prefix plus a run of the newest messages, no leading reply
    after a cut, and maximal."""
    rng = random.Random(0)
    ok = checked_maximal = replies_before_cut = 0
    for _ in range(400):
        n_sys = rng.choice([0, 1, 2])
        chat = [msg("system", " ".join("s" * rng.randint(1, 4))) for _ in range(n_sys)]
        chat += [msg(rng.choice(["user", "assistant"]), " ".join("w" * rng.randint(1, 9)))
                 for _ in range(rng.randint(1, 15))]
        limit, extra = rng.randint(5, 60), rng.randint(0, 3)
        try:
            kept = truncate(chat, limit, per_message=extra)
        except ValueError:
            continue
        ok += 1
        assert total_tokens(kept, per_message=extra) <= limit
        rest = chat[n_sys:]
        n_kept = len(kept) - n_sys
        assert kept[:n_sys] == chat[:n_sys]
        assert kept[n_sys:] == rest[len(rest) - n_kept:]                      # a run of the newest
        if n_kept < len(rest):
            assert kept[n_sys]["role"] != "assistant"                         # a cut never leaves a reply first
            older = rest[len(rest) - n_kept - 1]
            if older["role"] == "assistant":
                replies_before_cut += 1                                       # may have been dropped as an orphan
            else:                                                             # maximal: the next older one cannot fit
                checked_maximal += 1
                assert total_tokens(kept, per_message=extra) + count_tokens(older["content"]) + extra > limit
    assert ok > 100 and checked_maximal > 20 and replies_before_cut > 20      # the checks above were not vacuous


def test_rereading_history_grows_with_the_square_of_the_chat():
    """tokens_read must equal tokens_per_message x n^2, so a 10x longer chat re-reads 100x as much."""
    assert tokens_read(3, 10) == 10 + 30 + 50 == 90      # by hand: 1, 3, 5 messages of 10 tokens
    for n in (1, 2, 10, 20):
        assert tokens_read(n, 50) == 50 * n * n
    assert tokens_read(10, 50) == 5000 and tokens_read(20, 50) == 20000


def test_a_window_stops_the_growth():
    # window 30, messages of 10: replies read 10, 30, then 30 and 30 again (50 and 70 are cut to 30)
    assert tokens_read(3, 10, window=30) == 10 + 30 + 30 == 70
    n = 40
    assert tokens_read(n, 50, window=450) == 50 * (1 + 3 + 5 + 7 + 9) + 450 * (n - 5) == 17_000
    assert tokens_read(n, 50) == 80_000
    # lesson's last check-yourself question: 12 replies of 40-token messages
    assert tokens_read(12, 40) == 5_760
    assert tokens_read(12, 40, window=200) == 40 + 120 + 200 + 200 * 9 == 2_160


def test_a_window_of_eight_messages_sends_seven_because_of_the_orphan_reply():
    """400 tokens hold 8 messages of 50, but the oldest of the 8 is a reply whose question was cut."""
    assert tokens_read(10, 50, window=400) == 50 * (1 + 3 + 5 + 7) + 350 * 6 == 2_900


def test_attention_pairs_by_hand():
    assert attention_pairs(4) == 1 + 2 + 3 + 4 == 10
    assert attention_pairs(1000) == 500_500
    assert attention_pairs(2000) == 2_001_000            # doubling the prompt: about 4x the pairs