"""Tests for ExecutionQueue."""

import asyncio
import pytest

from skop_runner.execution_queue import ExecutionQueue, ExecutionCancelledError


@pytest.mark.asyncio
async def test_sequential_execution():
    """Submit 3 items to the same kernel, verify they execute in order."""
    queue = ExecutionQueue()
    order: list[int] = []

    async def make_fn(n: int):
        async def fn():
            order.append(n)
            await asyncio.sleep(0.01)
            return {"n": n}
        return fn

    coros = [queue.submit("k1", [f"cell-{i}"], await make_fn(i)) for i in range(3)]
    results = await asyncio.gather(*coros)

    assert order == [0, 1, 2]
    assert [r["n"] for r in results] == [0, 1, 2]
    await queue.cleanup_kernel("k1")


@pytest.mark.asyncio
async def test_cancel_queue_clears_pending():
    """Submit 3 items, cancel after first starts, verify 2-3 cancelled."""
    queue = ExecutionQueue()
    started = asyncio.Event()
    unblock = asyncio.Event()

    async def slow_fn():
        started.set()
        await unblock.wait()  # Wait until interrupt unblocks us
        return {"status": "interrupted"}

    async def fast_fn():
        return {"status": "done"}

    # Submit 3 items — first is slow, others are fast
    t1 = asyncio.create_task(queue.submit("k1", ["cell-1"], slow_fn))
    await started.wait()

    t2 = asyncio.create_task(queue.submit("k1", ["cell-2"], fast_fn))
    t3 = asyncio.create_task(queue.submit("k1", ["cell-3"], fast_fn))
    await asyncio.sleep(0.05)

    # Cancel — should cancel cell-2 and cell-3, interrupt cell-1
    async def mock_interrupt(kernel_id):
        unblock.set()  # Simulate interrupt unblocking the running execution

    result = await queue.cancel_queue("k1", interrupt_fn=mock_interrupt)

    assert set(result["cancelled_cells"]) == {"cell-2", "cell-3"}
    assert result["interrupted_cells"] == ["cell-1"]

    # t2 and t3 should raise ExecutionCancelledError
    with pytest.raises(ExecutionCancelledError):
        await t2
    with pytest.raises(ExecutionCancelledError):
        await t3

    # t1 should complete (interrupt unblocked it)
    r1 = await t1
    assert r1["status"] == "interrupted"

    await queue.cleanup_kernel("k1")


@pytest.mark.asyncio
async def test_cancel_queue_calls_interrupt():
    """Verify interrupt_fn is called for the running item."""
    queue = ExecutionQueue()
    started = asyncio.Event()
    unblock = asyncio.Event()

    async def slow_fn():
        started.set()
        await unblock.wait()
        return {}

    task = asyncio.create_task(queue.submit("k1", ["cell-1"], slow_fn))
    await started.wait()

    interrupted_kernel = None

    async def mock_interrupt(kernel_id):
        nonlocal interrupted_kernel
        interrupted_kernel = kernel_id
        unblock.set()

    await queue.cancel_queue("k1", interrupt_fn=mock_interrupt)
    assert interrupted_kernel == "k1"

    await task  # Should complete now
    await queue.cleanup_kernel("k1")


@pytest.mark.asyncio
async def test_get_queue_state():
    """Submit items, verify state shows queued + running."""
    queue = ExecutionQueue()
    started = asyncio.Event()
    unblock = asyncio.Event()

    async def slow_fn():
        started.set()
        await unblock.wait()
        return {}

    async def fast_fn():
        return {}

    task1 = asyncio.create_task(queue.submit("k1", ["cell-1"], slow_fn))
    await started.wait()

    # Queue more items
    task2 = asyncio.create_task(queue.submit("k1", ["cell-2"], fast_fn))
    task3 = asyncio.create_task(queue.submit("k1", ["cell-3"], fast_fn))
    await asyncio.sleep(0.05)

    state = queue.get_queue_state("k1")
    assert state["k1"]["running"] == ["cell-1"]
    assert state["k1"]["queued"] == 2

    # Cleanup
    unblock.set()
    await asyncio.gather(task1, task2, task3)
    await queue.cleanup_kernel("k1")


@pytest.mark.asyncio
async def test_different_kernels_parallel():
    """Submit to K1 and K2, verify parallel execution."""
    queue = ExecutionQueue()
    execution_log: list[tuple[str, str]] = []

    async def make_fn(kernel_id: str):
        async def fn():
            execution_log.append((kernel_id, "start"))
            await asyncio.sleep(0.05)
            execution_log.append((kernel_id, "end"))
            return {"kernel_id": kernel_id}
        return fn

    results = await asyncio.gather(
        queue.submit("k1", ["cell-a"], await make_fn("k1")),
        queue.submit("k2", ["cell-b"], await make_fn("k2")),
    )

    # Both should have started before either finished (parallel)
    starts = [i for i, (_, phase) in enumerate(execution_log) if phase == "start"]
    ends = [i for i, (_, phase) in enumerate(execution_log) if phase == "end"]
    assert max(starts) < min(ends)

    assert len(results) == 2
    await queue.cleanup_kernel("k1")
    await queue.cleanup_kernel("k2")


@pytest.mark.asyncio
async def test_cleanup_kernel():
    """Cleanup kernel, verify worker cancelled and queue empty."""
    queue = ExecutionQueue()
    started = asyncio.Event()

    async def slow_fn():
        started.set()
        await asyncio.sleep(10)
        return {}

    task = asyncio.create_task(queue.submit("k1", ["cell-1"], slow_fn))
    await started.wait()

    # Add a pending item
    task2 = asyncio.create_task(queue.submit("k1", ["cell-2"], slow_fn))
    await asyncio.sleep(0.05)

    await queue.cleanup_kernel("k1")

    # Worker should be gone
    assert "k1" not in queue._workers or queue._workers.get("k1", None) is None
    assert "k1" not in queue._queues
    assert "k1" not in queue._current

    # Pending future should be rejected
    with pytest.raises(ExecutionCancelledError):
        await task2

    # Running task's future is also rejected by cleanup
    with pytest.raises(ExecutionCancelledError):
        await task


@pytest.mark.asyncio
async def test_submit_after_cleanup_creates_new_worker():
    """After cleanup, a new submit should create a fresh worker."""
    queue = ExecutionQueue()

    async def fn():
        return {"value": 42}

    result = await queue.submit("k1", ["cell-1"], fn)
    assert result["value"] == 42
    await queue.cleanup_kernel("k1")

    # Submit again — should create a new worker
    result = await queue.submit("k1", ["cell-2"], fn)
    assert result["value"] == 42
    await queue.cleanup_kernel("k1")


@pytest.mark.asyncio
async def test_cancel_empty_queue():
    """Cancel with no queue should be a no-op."""
    queue = ExecutionQueue()

    result = await queue.cancel_queue("nonexistent", interrupt_fn=None)

    assert result["cancelled_cells"] == []
    assert result["interrupted_cells"] == []
