import json
from pathlib import Path

import pytest

from dlb.io import (
    SampleValidationError,
    expected_conditional_schedule,
    read_conditional_samples,
    validate_conditional_samples,
    write_conditional_samples_atomic,
)


def conditional_records(
    schedule: list[tuple[int, int]], *, sequence_length: int = 128
) -> list[dict[str, object]]:
    """Return independently constructed records for a supplied publication schedule."""

    records: list[dict[str, object]] = []
    for sample_id, (prompt_id, completion_id) in enumerate(schedule):
        prefix = [sample_id % 10] * 64
        continuation = [(sample_id + 1) % 10] * (sequence_length - 64)
        records.append(
            {
                "sample_id": sample_id,
                "prompt_id": prompt_id,
                "completion_id": completion_id,
                "source_index": sample_id + 10,
                "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,
            }
        )
    return records


def test_schedule_lists_quality_before_diversity_completions() -> None:
    """Catch a schedule that changes the public prompt/completion ordering."""

    assert expected_conditional_schedule(3, 2, 3) == [
        (0, 0),
        (1, 0),
        (2, 0),
        (0, 1),
        (1, 1),
        (0, 2),
        (1, 2),
    ]


def test_writer_and_reader_round_trip_a_parameterized_lm1b_artifact(tmp_path: Path) -> None:
    """Catch conditional I/O that loses validated records or small smoke schedules."""

    schedule = expected_conditional_schedule(2, 1, 2)
    path = tmp_path / "conditional.jsonl"

    assert write_conditional_samples_atomic(
        path,
        conditional_records(schedule),
        expected=3,
        schedule=schedule,
        sequence_length=128,
        vocab_size=10,
    ) == 3
    assert validate_conditional_samples(
        path,
        expected=3,
        schedule=schedule,
        sequence_length=128,
        vocab_size=10,
    ) == 3
    assert [(record.prompt_id, record.completion_id) for record in read_conditional_samples(path)] == schedule


@pytest.mark.parametrize(
    ("mutation", "message"),
    [
        (
            lambda records: records[0].update(
                continuation_token_ids=[1] * 63,
                full_token_ids=records[0]["prefix_token_ids"] + [1] * 63,
            ),
            "sequence_length",
        ),
        (
            lambda records: records[0].update(
                continuation_token_ids=[10] + [1] * 63,
                full_token_ids=records[0]["prefix_token_ids"] + [10] + [1] * 63,
            ),
            "vocabulary",
        ),
    ],
)
def test_validator_rejects_wrong_model_length_or_vocabulary_token(
    tmp_path: Path, mutation, message: str
) -> None:
    """Catch samples that cannot be evaluated on the declared dataset model contract."""

    schedule = [(0, 0)]
    records = conditional_records(schedule)
    mutation(records)
    path = tmp_path / "conditional.jsonl"
    path.write_text("\n".join(json.dumps(record) for record in records) + "\n", encoding="utf-8")

    with pytest.raises(SampleValidationError, match=message):
        validate_conditional_samples(
            path,
            expected=1,
            schedule=schedule,
            sequence_length=128,
            vocab_size=10,
        )


def test_reader_rejects_duplicate_json_keys(tmp_path: Path) -> None:
    """Catch ambiguous JSON that could hide a changed conditional field."""

    path = tmp_path / "conditional.jsonl"
    record = conditional_records([(0, 0)])[0]
    payload = json.dumps(record, separators=(",", ":")).replace(
        '"prompt_id":0,', '"prompt_id":0,"prompt_id":1,', 1
    )
    path.write_text(payload + "\n", encoding="utf-8")

    with pytest.raises(SampleValidationError, match="duplicate JSON key"):
        list(read_conditional_samples(path))


def test_writer_rejects_duplicate_schedule_entry_without_replacing_final(tmp_path: Path) -> None:
    """Catch publication that replaces a finished shard before schedule validation."""

    final = tmp_path / "samples.jsonl"
    final.write_text("old\n", encoding="utf-8")
    records = conditional_records(expected_conditional_schedule())
    records[1024]["prompt_id"] = 1

    with pytest.raises(SampleValidationError, match="expected prompt/completion"):
        write_conditional_samples_atomic(
            final,
            records,
            expected=2048,
            sequence_length=128,
            vocab_size=30_522,
        )

    assert final.read_text(encoding="utf-8") == "old\n"


def test_validator_requires_all_schedule_records_when_expected_is_none(tmp_path: Path) -> None:
    """Catch a partial conditional shard accepted when its caller omits expected."""

    path = tmp_path / "conditional.jsonl"
    path.write_text(json.dumps(conditional_records([(0, 0)])[0]) + "\n", encoding="utf-8")

    with pytest.raises(SampleValidationError, match="expected 2 records, found 1"):
        validate_conditional_samples(
            path,
            expected=None,
            schedule=[(0, 0), (1, 0)],
            sequence_length=128,
            vocab_size=10,
        )


@pytest.mark.parametrize("operation", ["validate", "write"])
def test_conditional_publication_requires_explicit_dataset_contract(
    tmp_path: Path, operation: str
) -> None:
    """Catch accidental LM1B defaults used to validate an unspecified dataset artifact."""

    path = tmp_path / "conditional.jsonl"
    records = conditional_records([(0, 0)])
    path.write_text(json.dumps(records[0]) + "\n", encoding="utf-8")

    with pytest.raises(TypeError):
        if operation == "validate":
            validate_conditional_samples(path, expected=1, schedule=[(0, 0)])
        else:
            write_conditional_samples_atomic(path, records, expected=1, schedule=[(0, 0)])


@pytest.mark.parametrize(
    "schedule",
    [
        [(0, 0), (0, 0)],
        [(1024, 0)],
        [(0, 5)],
        [(0, 0), (1, 0), (0, 1), (1, 1), (0, 2)],
    ],
)
def test_writer_rejects_noncanonical_custom_schedule(tmp_path: Path, schedule) -> None:
    """Catch duplicate, out-of-range, or incomplete prompt/completion schedule definitions."""

    with pytest.raises(ValueError, match="conditional schedule"):
        write_conditional_samples_atomic(
            tmp_path / "conditional.jsonl",
            [],
            expected=None,
            schedule=schedule,
            sequence_length=128,
            vocab_size=10,
        )