"""Tests for cancellation token behavior."""

import asyncio
from concurrent.futures import Future

import pytest

from agentbyte.cancellation_token import CancellationToken


def test_cancel_sets_state() -> None:
    token = CancellationToken()
    assert token.is_cancelled() is False

    token.cancel()

    assert token.is_cancelled() is True


def test_add_callback_invoked_on_cancel() -> None:
    token = CancellationToken()
    called = {"value": False}

    def callback() -> None:
        called["value"] = True

    token.add_callback(callback)
    token.cancel()

    assert called["value"] is True


def test_add_callback_after_cancel_invokes_immediately() -> None:
    token = CancellationToken()
    token.cancel()
    called = {"value": False}

    def callback() -> None:
        called["value"] = True

    token.add_callback(callback)

    assert called["value"] is True


def test_link_future_cancelled_when_token_cancelled() -> None:
    token = CancellationToken()
    future: Future[None] = Future()

    linked = token.link_future(future)
    token.cancel()

    assert linked.cancelled() is True


@pytest.mark.asyncio
async def test_link_async_task_cancelled_when_token_cancelled() -> None:
    token = CancellationToken()

    async def sleeper() -> None:
        await asyncio.sleep(10)

    task = asyncio.create_task(sleeper())
    token.link_future(task)
    token.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task
