from decimal import Decimal

import pytest

from cost import common_prefix_len, decode_positions_read, request_cost, sequential_passes


def test_cost_by_hand():
    """1,000 input tokens at 3 dollars per million is 0.003 dollars. 2,000 at 3 plus 500 out at 15
    is 0.006 + 0.0075 = 0.0135."""
    assert request_cost(1000, 0, 0, "3", "0.3", "15") == Decimal("0.003")
    assert request_cost(2000, 0, 500, "3", "0.3", "15") == Decimal("0.0135")


def test_cached_tokens_use_their_own_rate():
    assert request_cost(0, 1_000_000, 0, 3, "0.3", 15) == Decimal("0.3")
    assert request_cost(500, 1500, 400, 2, "0.2", 10) == Decimal("0.0053")


def test_check_yourself_values():
    """The quiz answers in the lesson: 2,000 x 2 + 6,000 x 0.2 + 500 x 10 = 4,000 + 1,200 + 5,000
    = 10,200 millionths of a dollar; with no cache 8,000 x 2 + 5,000 = 21,000."""
    assert request_cost(2000, 6000, 500, 2, "0.2", 10) == Decimal("0.0102")
    assert request_cost(8000, 0, 500, 2, "0.2", 10) == Decimal("0.021")
    assert decode_positions_read(1000, 1000) == 999 * 1000 + 999 * 1000 // 2 == 1_498_500
    assert decode_positions_read(2000, 1000) == 2_497_500
    assert sequential_passes(1000, 1000) == sequential_passes(2000, 1000) == 1000


def test_money_is_exact_where_float_is_not():
    """Guards a float sum: 0.1 + 0.2 is not 0.3 in binary, and a bill should not carry that."""
    assert request_cost(1_000_000, 1_000_000, 0, "0.1", "0.2", 0) == Decimal("0.3")


@pytest.mark.parametrize("counts, prices", [((-1, 0, 0), (1, 1, 1)), ((0, 0, 0), (1, -1, 1))])
def test_negative_input_is_refused(counts, prices):
    with pytest.raises(ValueError):
        request_cost(*counts, *prices)


def test_passes_one_per_output_token_and_the_prompt_is_one_pass():
    assert sequential_passes(2000, 500) == 500
    assert sequential_passes(5, 1) == 1
    assert sequential_passes(5, 0) == 1                   # the prompt pass still happens
    assert sequential_passes(50_000, 3) == sequential_passes(5, 3)   # prompt length does not add passes


def test_decode_positions_by_hand():
    """3 prompt tokens, 4 output tokens: decode passes 1, 2, 3 read 3+1, 3+2, 3+3 = 4 + 5 + 6 = 15."""
    assert decode_positions_read(3, 4) == 15
    assert decode_positions_read(2000, 500) == 499 * 2000 + 499 * 500 // 2 == 1_122_750
    assert decode_positions_read(7, 0) == decode_positions_read(7, 1) == 0


def test_decode_positions_grow_faster_than_output_length():
    """Doubling the output more than doubles the reading: the reason long outputs are costly."""
    assert decode_positions_read(100, 2000) > 2 * decode_positions_read(100, 1000)


def test_common_prefix_len():
    assert common_prefix_len([1, 2, 3, 4], [1, 2, 9, 4]) == 2
    assert common_prefix_len([1, 2], [1, 2, 3]) == 2
    assert common_prefix_len([], [1]) == 0