# -*- coding: utf-8 -*-
"""Location: ./tests/unit/mcpgateway/services/test_cancellation_service.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Tests for cancellation_service.
"""

# Standard
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock

# Third-Party
import pytest

# First-Party
from mcpgateway.services.cancellation_service import CancellationService


@pytest.mark.asyncio
async def test_initialize_with_redis(monkeypatch):
    service = CancellationService()
    monkeypatch.setattr(service, "_listen_for_cancellations", AsyncMock())
    monkeypatch.setattr("mcpgateway.services.cancellation_service.get_redis_client", AsyncMock(return_value=MagicMock()))

    created_tasks = []

    def _fake_create_task(coro):
        created_tasks.append(coro)
        task = MagicMock()
        task.done.return_value = False
        return task

    monkeypatch.setattr("mcpgateway.services.cancellation_service.asyncio.create_task", _fake_create_task)

    await service.initialize()
    assert service._initialized is True
    assert service._pubsub_task is not None
    assert created_tasks


@pytest.mark.asyncio
async def test_initialize_is_idempotent(monkeypatch):
    service = CancellationService()
    mock_get = AsyncMock(return_value=None)
    monkeypatch.setattr("mcpgateway.services.cancellation_service.get_redis_client", mock_get)

    await service.initialize()
    await service.initialize()

    # Second call should return early without re-fetching.
    assert mock_get.await_count == 1


@pytest.mark.asyncio
async def test_initialize_logs_warning_on_exception(monkeypatch):
    service = CancellationService()
    monkeypatch.setattr("mcpgateway.services.cancellation_service.get_redis_client", AsyncMock(side_effect=RuntimeError("boom")))

    await service.initialize()
    assert service._initialized is True


@pytest.mark.asyncio
async def test_shutdown_cancels_task():
    service = CancellationService()

    class DummyTask:
        def __init__(self):
            self.cancel_called = False

        def done(self):
            return False

        def cancel(self):
            self.cancel_called = True

        def __await__(self):
            async def _noop():
                return None

            return _noop().__await__()

    task = DummyTask()
    service._pubsub_task = task

    await service.shutdown()
    assert task.cancel_called is True


@pytest.mark.asyncio
async def test_shutdown_handles_cancelled_error_from_task():
    service = CancellationService()

    class DummyTask:
        def done(self):
            return False

        def cancel(self):
            return None

        def __await__(self):
            async def _raise():
                raise asyncio.CancelledError()

            return _raise().__await__()

    service._pubsub_task = DummyTask()
    await service.shutdown()


@pytest.mark.asyncio
async def test_cancel_run_unknown_publishes(monkeypatch):
    service = CancellationService()
    service._publish_cancellation = AsyncMock()

    result = await service.cancel_run("missing", reason="test")
    assert result is False
    service._publish_cancellation.assert_awaited_once_with("missing", "test")


@pytest.mark.asyncio
async def test_cancel_run_known_executes_callback(monkeypatch):
    service = CancellationService()
    callback = AsyncMock()
    await service.register_run("run-1", name="tool", cancel_callback=callback)
    service._publish_cancellation = AsyncMock()

    result = await service.cancel_run("run-1", reason="stop")
    assert result is True
    callback.assert_awaited_once_with("stop")
    service._publish_cancellation.assert_awaited_once_with("run-1", "stop")


@pytest.mark.asyncio
async def test_cancel_run_local_handles_callback_error():
    service = CancellationService()

    async def _boom(_reason):
        raise RuntimeError("bad")

    await service.register_run("run-1", cancel_callback=_boom)
    result = await service._cancel_run_local("run-1", reason="x")
    assert result is True


@pytest.mark.asyncio
async def test_cancel_run_local_run_not_found():
    service = CancellationService()
    assert await service._cancel_run_local("missing") is False


@pytest.mark.asyncio
async def test_cancel_run_local_already_cancelled():
    service = CancellationService()
    await service.register_run("run-1")
    service._runs["run-1"]["cancelled"] = True
    assert await service._cancel_run_local("run-1") is True


@pytest.mark.asyncio
async def test_cancel_run_local_callback_success():
    service = CancellationService()
    callback = AsyncMock()
    await service.register_run("run-1", cancel_callback=callback)
    assert await service._cancel_run_local("run-1", reason="ok") is True
    callback.assert_awaited_once_with("ok")


@pytest.mark.asyncio
async def test_publish_cancellation_no_redis():
    service = CancellationService()
    service._redis = None
    await service._publish_cancellation("run-1", reason="no-redis")


@pytest.mark.asyncio
async def test_publish_cancellation_redis_error():
    service = CancellationService()
    redis = AsyncMock()
    redis.publish.side_effect = RuntimeError("fail")
    service._redis = redis
    await service._publish_cancellation("run-1", reason="boom")


@pytest.mark.asyncio
async def test_publish_cancellation_success_logs_debug(monkeypatch):
    service = CancellationService()
    redis = AsyncMock()
    service._redis = redis

    await service._publish_cancellation("run-1", reason="ok")
    redis.publish.assert_awaited_once()


@pytest.mark.asyncio
async def test_get_status_and_is_registered():
    service = CancellationService()
    await service.register_run("run-1", name="tool")
    assert await service.is_registered("run-1") is True
    status = await service.get_status("run-1")
    assert status["name"] == "tool"
    await service.unregister_run("run-1")
    assert await service.is_registered("run-1") is False


@pytest.mark.asyncio
async def test_cancel_run_already_cancelled(monkeypatch):
    service = CancellationService()
    await service.register_run("run-1", name="tool")
    service._runs["run-1"]["cancelled"] = True
    service._publish_cancellation = AsyncMock()

    result = await service.cancel_run("run-1", reason="again")
    assert result is True
    service._publish_cancellation.assert_not_awaited()


@pytest.mark.asyncio
async def test_listen_for_cancellations_returns_when_not_initialized():
    service = CancellationService()
    assert await service._listen_for_cancellations() is None


@pytest.mark.asyncio
async def test_listen_for_cancellations_timeout_continues(monkeypatch):
    service = CancellationService()

    class DummyPubSub:
        async def subscribe(self, _channel):
            return None

        async def unsubscribe(self, _channel):
            return None

        async def aclose(self):
            return None

        async def get_message(self, **_kwargs):
            return None

    class DummyRedis:
        def pubsub(self):
            return DummyPubSub()

    service._redis = DummyRedis()

    calls = {"n": 0}

    async def _fake_wait_for(*_args, **_kwargs):
        # Ensure we don't leak an un-awaited coroutine created by pubsub.get_message(...).
        if _args:
            aw = _args[0]
            close = getattr(aw, "close", None)
            if callable(close):
                close()
        calls["n"] += 1
        if calls["n"] == 1:
            raise asyncio.TimeoutError()
        raise asyncio.CancelledError()

    monkeypatch.setattr("mcpgateway.services.cancellation_service.asyncio.wait_for", _fake_wait_for)

    with pytest.raises(asyncio.CancelledError):
        await service._listen_for_cancellations()


@pytest.mark.asyncio
async def test_listen_for_cancellations_skips_null_run_id(monkeypatch):
    service = CancellationService()

    messages = [
        {"type": "message", "data": json.dumps({"run_id": None, "reason": "stop"}).encode()},
    ]

    class DummyPubSub:
        def __init__(self, items):
            self._items = list(items)
            self.unsubscribe_called = False

        async def subscribe(self, _channel):
            return None

        async def get_message(self, **_kwargs):
            if self._items:
                return self._items.pop(0)
            raise asyncio.CancelledError()

        async def unsubscribe(self, _channel):
            self.unsubscribe_called = True

        async def aclose(self):
            return None

    class DummyRedis:
        def __init__(self, pubsub):
            self._pubsub = pubsub

        def pubsub(self):
            return self._pubsub

    pubsub = DummyPubSub(messages)
    service._redis = DummyRedis(pubsub)
    monkeypatch.setattr(service, "_cancel_run_local", AsyncMock())

    task = asyncio.create_task(service._listen_for_cancellations())
    with pytest.raises(asyncio.CancelledError):
        await task

    service._cancel_run_local.assert_not_awaited()
    assert pubsub.unsubscribe_called is True


@pytest.mark.asyncio
async def test_listen_for_cancellations_pubsub_error_is_caught():
    service = CancellationService()

    class DummyRedis:
        def pubsub(self):
            raise RuntimeError("boom")

    service._redis = DummyRedis()
    # Should not raise.
    assert await service._listen_for_cancellations() is None


@pytest.mark.asyncio
async def test_listen_for_cancellations_cleanup_uses_close_when_aclose_missing():
    service = CancellationService()

    class DummyPubSub:
        def __init__(self):
            self.closed = False

        async def subscribe(self, _channel):
            return None

        async def get_message(self, **_kwargs):
            raise asyncio.CancelledError()

        async def unsubscribe(self, _channel):
            return None

        async def close(self):
            self.closed = True

    class DummyRedis:
        def __init__(self, pubsub):
            self._pubsub = pubsub

        def pubsub(self):
            return self._pubsub

    pubsub = DummyPubSub()
    service._redis = DummyRedis(pubsub)

    task = asyncio.create_task(service._listen_for_cancellations())
    with pytest.raises(asyncio.CancelledError):
        await task

    assert pubsub.closed is True


@pytest.mark.asyncio
async def test_listen_for_cancellations_cleanup_exception_is_swallowed(monkeypatch):
    service = CancellationService()

    class DummyPubSub:
        async def subscribe(self, _channel):
            return None

        async def get_message(self, **_kwargs):
            raise asyncio.CancelledError()

        async def unsubscribe(self, _channel):
            raise RuntimeError("boom")

        async def aclose(self):
            return None

    class DummyRedis:
        def __init__(self, pubsub):
            self._pubsub = pubsub

        def pubsub(self):
            return self._pubsub

    service._redis = DummyRedis(DummyPubSub())
    debug = MagicMock()
    monkeypatch.setattr("mcpgateway.services.cancellation_service.logger.debug", debug)

    task = asyncio.create_task(service._listen_for_cancellations())
    with pytest.raises(asyncio.CancelledError):
        await task

    assert debug.called


@pytest.mark.asyncio
async def test_listen_for_cancellations_processes_messages(monkeypatch):
    service = CancellationService()

    messages = [
        None,
        {"type": "subscribe"},
        {"type": "message", "data": json.dumps({"run_id": "run-1", "reason": "stop"}).encode()},
        {"type": "message", "data": b"not-json"},
    ]

    class DummyPubSub:
        def __init__(self, items):
            self._items = list(items)
            self.unsubscribe_called = False
            self.closed = False

        async def subscribe(self, _channel):
            return None

        async def get_message(self, **_kwargs):
            if self._items:
                return self._items.pop(0)
            raise asyncio.CancelledError()

        async def unsubscribe(self, _channel):
            self.unsubscribe_called = True

        async def aclose(self):
            self.closed = True

    class DummyRedis:
        def __init__(self, pubsub):
            self._pubsub = pubsub

        def pubsub(self):
            return self._pubsub

    pubsub = DummyPubSub(messages)
    service._redis = DummyRedis(pubsub)
    monkeypatch.setattr(service, "_cancel_run_local", AsyncMock())

    task = asyncio.create_task(service._listen_for_cancellations())
    with pytest.raises(asyncio.CancelledError):
        await task

    service._cancel_run_local.assert_awaited_with("run-1", reason="stop")
    assert pubsub.unsubscribe_called is True
    assert pubsub.closed is True
