from __future__ import annotations

import asyncio

import pytest

from backplane_runtime.backends.base import BackendInput
from backplane_runtime.backends.mock import EvidenceBackend
from backplane_runtime.scheduler import BatchScheduler
from backplane_runtime.errors import RequestLimitError


def item(index: int, stack: str = "same") -> BackendInput:
    return BackendInput(
        prompt=f"request {index}",
        max_new_tokens=16,
        tenant_id="tenant",
        request_id=f"req-{index}",
        stack_key=stack,
    )


def test_compatible_requests_microbatch():
    async def scenario():
        scheduler = BatchScheduler(
            EvidenceBackend(), max_batch_size=8, batch_window_ms=20
        )
        results = await asyncio.gather(*[scheduler.submit(item(i)) for i in range(6)])
        assert {batch_size for _, _, batch_size in results} == {6}
        assert all(result.metadata["batch_size"] == 6 for result, _, _ in results)
        await scheduler.close()

    asyncio.run(scenario())


def test_incompatible_stacks_are_separated():
    async def scenario():
        scheduler = BatchScheduler(
            EvidenceBackend(), max_batch_size=8, batch_window_ms=20
        )
        results = await asyncio.gather(
            scheduler.submit(item(1, "a")),
            scheduler.submit(item(2, "b")),
            scheduler.submit(item(3, "a")),
        )
        batch_sizes = [batch_size for _, _, batch_size in results]
        assert sorted(batch_sizes) == [1, 2, 2]
        await scheduler.close()

    asyncio.run(scenario())


def test_queue_measurement_excludes_backend_execution(monkeypatch):
    async def scenario():
        now=[0.0]
        monkeypatch.setattr('backplane_runtime.scheduler.time.perf_counter',lambda:now[0])
        class SlowClockBackend(EvidenceBackend):
            async def generate_batch(self,items):
                now[0]=10.0
                return await super().generate_batch(items)
        scheduler=BatchScheduler(SlowClockBackend(),batch_window_ms=0)
        _,queued,_=await scheduler.submit(item(1))
        assert queued==0.0
        await scheduler.close()
    asyncio.run(scenario())


def test_close_rejects_new_work_and_survives_caller_cancellation():
    async def scenario():
        started=asyncio.Event();finish=asyncio.Event()
        class Blocking(EvidenceBackend):
            async def generate_batch(self,items):
                started.set();await finish.wait();return await super().generate_batch(items)
        scheduler=BatchScheduler(Blocking(),batch_window_ms=0)
        pending=asyncio.create_task(scheduler.submit(item(1)));await started.wait()
        closing=asyncio.create_task(scheduler.close());await asyncio.sleep(0)
        with pytest.raises(RuntimeError,match='closed'):await scheduler.submit(item(2))
        closing.cancel()
        with pytest.raises(asyncio.CancelledError):await closing
        finish.set();await asyncio.wait_for(pending,1)
        await asyncio.wait_for(scheduler.close(),1)
    asyncio.run(scenario())


def test_pending_queue_is_bounded_and_cancelled_work_is_not_executed():
    async def scenario():
        started=asyncio.Event();finish=asyncio.Event();seen=[]
        class Blocking(EvidenceBackend):
            async def generate_batch(self,items):
                seen.extend(entry.request_id for entry in items);started.set();await finish.wait()
                return await super().generate_batch(items)
        scheduler=BatchScheduler(Blocking(),max_batch_size=1,batch_window_ms=0,max_pending_requests=1)
        first=asyncio.create_task(scheduler.submit(item(1)));await started.wait()
        second=asyncio.create_task(scheduler.submit(item(2)));await asyncio.sleep(0)
        with pytest.raises(RequestLimitError,match='queue is full'):await scheduler.submit(item(3))
        second.cancel()
        with pytest.raises(asyncio.CancelledError):await second
        finish.set();await first;await scheduler.close()
        assert seen==['req-1']
    asyncio.run(scenario())


def test_backend_cancellation_settles_waiters_instead_of_hanging():
    async def scenario():
        class Cancelled(EvidenceBackend):
            async def generate_batch(self,items):raise asyncio.CancelledError()
        scheduler=BatchScheduler(Cancelled(),batch_window_ms=0)
        with pytest.raises(RuntimeError,match='cancelled'):
            await asyncio.wait_for(scheduler.submit(item(1)),1)
        with pytest.raises(RuntimeError,match='cancelled'):await scheduler.submit(item(2))
        await scheduler.close()
    asyncio.run(scenario())
