"""The batch runner: parsing the prompt window, scheduling, retry, error
rows, events and eviction. The runner holds no model — the caller hands it a
coroutine — so these run it for real on a stub answerer and assert on the
run's own accounting, which is the module's whole contract."""

import asyncio

import pytest

from quber.playground.batch import BatchRunner, parse_questions


def test_parse_questions_one_per_line_blanks_dropped_duplicates_kept():
    text = "What was Q3 revenue?\n\n  What drove margins?  \nWhat was Q3 revenue?\n"
    assert parse_questions(text) == [
        "What was Q3 revenue?",
        "What drove margins?",
        "What was Q3 revenue?",
    ]
    assert parse_questions("\n \n") == []


async def retrieved(question):
    return [question]


async def answered(question, context):
    return {"shape": "scalar", "answer": question, "value": {"value": "1"}, "references": []}


def wait_done(runner, job_id):
    async def drain():
        async for event in runner.events(job_id):
            if event.get("type") == "done":
                return event

    return drain()


def test_rows_land_per_question_and_state_reports_them():
    async def run():
        runner = BatchRunner(concurrency=2)
        state = runner.start(1, ["q1", "q2", "q3"], "value", retrieved, answered)
        assert state.total == 3 and state.running
        done = await wait_done(runner, state.job_id)
        assert done == {"type": "done", "total": 3, "done": 3, "failed": 0}
        final = runner.get(state.job_id)
        assert final is not None
        assert [r.question for r in final.rows] == ["q1", "q2", "q3"]
        first = final.rows[0].response
        assert first is not None and first["answer"] == "q1"
        assert not final.running

    asyncio.run(run())


def test_a_question_that_fails_becomes_a_row_carrying_its_error():
    async def flaky(question, context):
        if question == "bad":
            raise RuntimeError("boom")
        return await answered(question, context)

    async def run():
        runner = BatchRunner()
        state = runner.start(1, ["ok", "bad"], "auto", retrieved, flaky)
        await wait_done(runner, state.job_id)
        final = runner.get(state.job_id)
        assert final is not None
        assert final.failed == 1 and final.done == 2
        bad = final.rows[1]
        assert bad.error == "boom" and bad.response is None

    asyncio.run(run())


def test_retry_on_overload_then_success():
    class Overloaded(Exception):
        status_code = 529

    calls = {"n": 0}

    async def once_overloaded(question, context):
        calls["n"] += 1
        if calls["n"] == 1:
            raise Overloaded("overloaded")
        return await answered(question, context)

    async def run():
        runner = BatchRunner()
        # Shrink the first backoff so the test is quick.
        import quber.playground.batch as batch_module

        original = batch_module.FIRST_RETRY_DELAY
        batch_module.FIRST_RETRY_DELAY = 0.01
        try:
            state = runner.start(1, ["q"], "auto", retrieved, once_overloaded)
            await wait_done(runner, state.job_id)
        finally:
            batch_module.FIRST_RETRY_DELAY = original
        final = runner.get(state.job_id)
        assert final is not None
        assert final.failed == 0
        assert calls["n"] == 2

    asyncio.run(run())


def test_a_non_retryable_error_is_not_retried():
    calls = {"n": 0}

    async def always_bad(question, context):
        calls["n"] += 1
        raise ValueError("schema mismatch")

    async def run():
        runner = BatchRunner()
        state = runner.start(1, ["q"], "auto", retrieved, always_bad)
        await wait_done(runner, state.job_id)
        assert calls["n"] == 1
        final = runner.get(state.job_id)
        assert final is not None and final.failed == 1

    asyncio.run(run())


def test_eviction_past_the_cap_reports_the_run_as_gone():
    async def run():
        runner = BatchRunner(max_runs=2)
        first = runner.start(1, ["q"], "auto", retrieved, answered)
        second = runner.start(1, ["q"], "auto", retrieved, answered)
        third = runner.start(1, ["q"], "auto", retrieved, answered)
        for s in (second, third):
            await wait_done(runner, s.job_id)
        assert runner.get(first.job_id) is None
        assert runner.get(third.job_id) is not None

    asyncio.run(run())


def test_row_events_carry_index_total_and_the_row_itself():
    async def run():
        runner = BatchRunner(concurrency=1)
        state = runner.start(1, ["q1", "q2"], "auto", retrieved, answered)
        rows = []
        async for event in runner.events(state.job_id):
            if event["type"] == "row":
                rows.append(event)
            if event["type"] == "done":
                break
        assert [e["index"] for e in rows] == [0, 1]
        assert rows[0]["total"] == 2 and rows[0]["ok"] is True
        assert rows[0]["row"]["question"] == "q1"
        assert rows[0]["row"]["response"]["answer"] == "q1"

    asyncio.run(run())


def test_subscribing_to_a_finished_run_yields_done_immediately():
    async def run():
        runner = BatchRunner()
        state = runner.start(1, ["q"], "auto", retrieved, answered)
        await wait_done(runner, state.job_id)
        events = [e async for e in runner.events(state.job_id)]
        assert events == [{"type": "done", "total": 1, "done": 1, "failed": 0}]

    asyncio.run(run())


def test_starting_outside_a_loop_raises():
    runner = BatchRunner()
    with pytest.raises(RuntimeError):
        runner.start(1, ["q"], "auto", retrieved, answered)


def test_running_for_reports_only_in_flight_runs_for_that_document():
    async def run():
        runner = BatchRunner()
        state = runner.start(1, ["q"], "auto", retrieved, answered)
        await wait_done(runner, state.job_id)
        assert not runner.running_for(1)
        assert not runner.running_for(2)

    asyncio.run(run())