import pytest

from central.app.ai.batch_runner import run_batches
from central.app.ai.fake_provider import DeterministicFakeSemanticProvider
from central.app.ai.llm_provider import LlmEvalItem


def _item(work_item_id):
    return LlmEvalItem(
        work_item_id=work_item_id,
        description=work_item_id,
        normalized_description=work_item_id,
        brand=None,
        unit=None,
        segment="test",
        candidates=[("A", "candidate A", 1.0), ("B", "candidate B", 0.5)],
    )


@pytest.mark.parametrize("batch_size", [5, 10, 20])
def test_run_batches_correlates_shuffled_answers_by_id(batch_size):
    items = [_item(f"w{i}") for i in range(batch_size + 1)]
    provider = DeterministicFakeSemanticProvider(shuffle=True)

    result = run_batches(items, provider, batch_size=batch_size)

    assert set(result.results_by_work_item_id) == {item.work_item_id for item in items}
    assert all(result.results_by_work_item_id[item.work_item_id]["work_item_id"] == item.work_item_id for item in items)
    assert result.errors == {}
    assert result.provider_calls == 2


@pytest.mark.parametrize("scenario", ["MISSING_ITEM", "DUPLICATE_ITEM", "INVALID_JSON", "OUT_OF_CANDIDATE"])
def test_run_batches_preserves_failed_ids_for_contract_errors(scenario):
    items = [_item("w1"), _item("w2")]
    provider = DeterministicFakeSemanticProvider()

    result = run_batches(items, provider, batch_size=2, scenario=scenario)

    assert result.results_by_work_item_id == {}
    assert set(result.errors) == {"w1", "w2"}
    assert result.retries == 0
    assert result.provider_calls == 1


@pytest.mark.parametrize("scenario", ["TIMEOUT", "429", "500"])
def test_run_batches_retries_transport_errors_and_preserves_failed_ids(scenario):
    items = [_item("w1"), _item("w2")]
    provider = DeterministicFakeSemanticProvider()

    result = run_batches(items, provider, batch_size=2, max_retries=2, scenario=scenario)

    assert result.results_by_work_item_id == {}
    assert set(result.errors) == {"w1", "w2"}
    assert result.provider_calls == 3
    assert result.retries == 2


def test_run_batches_allows_partial_failure_by_item():
    items = [_item("w1"), _item("w2"), _item("w3")]
    provider = DeterministicFakeSemanticProvider(fail_work_item_ids={"w2"}, failure_scenario="TIMEOUT")

    result = run_batches(items, provider, batch_size=1, max_retries=1)

    assert set(result.results_by_work_item_id) == {"w1", "w3"}
    assert set(result.errors) == {"w2"}
    assert result.provider_calls == 4
    assert result.retries == 1