import pytest
from pydantic import ValidationError

from dlb.schema import ConditionalSampleRecord


def valid_conditional_record() -> dict[str, object]:
    """Return a hand-authored LM1B-sized conditional publication record."""

    prefix = [7] * 64
    continuation = [8] * 64
    return {
        "sample_id": 0,
        "prompt_id": 0,
        "completion_id": 0,
        "source_index": 12,
        "prefix_token_ids": prefix,
        "continuation_token_ids": continuation,
        "reference_token_ids": [9] * 64,
        "full_token_ids": prefix + continuation,
        "prefix_text": "fixed prefix",
        "continuation_text": "generated continuation",
        "reference_text": "source continuation",
        "full_text": "fixed prefix generated continuation",
        "seed": 42,
        "generation_seconds": 0.25,
        "prefix_exact_match": True,
    }


def test_conditional_record_binds_prefix_full_and_suffix_slices() -> None:
    """Catch publication records whose retained prefix or suffix is inconsistent."""

    record = ConditionalSampleRecord.model_validate(valid_conditional_record())

    assert record.full_token_ids[:64] == record.prefix_token_ids
    assert record.full_token_ids[64:] == record.continuation_token_ids
    assert record.prefix_exact_match is True


@pytest.mark.parametrize(
    ("mutation", "message"),
    [
        (lambda value: value.update(prefix_exact_match=False), "prefix_exact_match"),
        (lambda value: value["full_token_ids"].__setitem__(0, 999), "prefix"),
        (lambda value: value.update(reference_token_ids=[1] * 63), "64"),
    ],
)
def test_conditional_record_rejects_invalid_publication(mutation, message: str) -> None:
    """Catch loss of the hard-prefix or aligned-reference publication invariants."""

    value = valid_conditional_record()
    mutation(value)

    with pytest.raises(ValidationError, match=message):
        ConditionalSampleRecord.model_validate(value)