from embedding_chunking import CHUNK_OVERLAP, CHUNK_SIZE, ENCODER, chunk_text


def _text_of_n_tokens(n: int) -> str:
    """Build a string that encodes to exactly n cl100k tokens.

    Single ASCII letters encode to one token each, so n letters → n tokens.
    """
    token_ids = ENCODER.encode("a " * n)[:n]
    return ENCODER.decode(token_ids)


def test_short_text_returns_single_chunk():
    text = _text_of_n_tokens(100)
    chunks = chunk_text(text)
    assert len(chunks) == 1
    assert chunks[0].index == 0
    assert chunks[0].token_count == 100
    assert chunks[0].text == text


def test_1500_token_text_chunks_with_overlap():
    # Brief asserted 3 chunks for 1500 tokens, but with step=448 the
    # sliding-window starts are [0, 448, 896, 1344] → 4 chunks (three
    # full 512-token windows plus a 156-token tail). Keeping the
    # canonical sliding-window algorithm; updating the expectation.
    text = _text_of_n_tokens(1500)
    chunks = chunk_text(text)

    step = CHUNK_SIZE - CHUNK_OVERLAP
    assert len(chunks) == 4
    assert [c.index for c in chunks] == [0, 1, 2, 3]
    assert [c.token_count for c in chunks] == [
        CHUNK_SIZE,
        CHUNK_SIZE,
        CHUNK_SIZE,
        1500 - 3 * step,
    ]

    tokens = ENCODER.encode(text)
    for i, c in enumerate(chunks):
        start = i * step
        assert ENCODER.encode(c.text) == tokens[start : start + CHUNK_SIZE]


def test_600_token_text_returns_two_chunks():
    text = _text_of_n_tokens(600)
    chunks = chunk_text(text)

    assert len(chunks) == 2
    assert [c.index for c in chunks] == [0, 1]

    step = CHUNK_SIZE - CHUNK_OVERLAP
    assert chunks[0].token_count == CHUNK_SIZE
    assert chunks[1].token_count == 600 - step