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

Memory-backend unit tests for ``session_registry.py``.
This suite exercises the in-memory implementation of
:pyfile:`mcpgateway/cache/session_registry.py`.

Covered behaviours
------------------
* add_session / get_session / get_session_sync / remove_session
* broadcast -> respond for **dict**, **list**, **str** payloads
* generate_response branches:
  - initialize (result + notifications)
  - ping
  - tools/list (with stubbed service + DB)
* handle_initialize_logic success, and both error branches


Includes comprehensive backend testing, error scenarios, and cleanup tasks.
"""

# Future
from __future__ import annotations

# Standard
import asyncio
import importlib
import json
import logging
import re
import sys
from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock, Mock, patch

# Third-Party
from fastapi import HTTPException
import orjson
import pytest

# First-Party
from mcpgateway.cache.session_registry import SessionMessageRecord, SessionRegistry
from mcpgateway.config import settings


# --------------------------------------------------------------------------- #
# Minimal SSE transport stub                                                  #
# --------------------------------------------------------------------------- #
class FakeSSETransport:
    """Stub implementing just the subset of the API used by SessionRegistry."""

    def __init__(self, session_id: str, connected: bool = True):
        self.session_id = session_id
        self._connected = connected
        self.sent: List[Any] = []
        self.disconnect_called = False

    async def disconnect(self) -> None:  # noqa: D401
        self._connected = False
        self.disconnect_called = True

    async def is_connected(self) -> bool:  # noqa: D401
        return self._connected

    async def send_message(self, msg) -> None:  # noqa: D401
        if not self._connected:
            raise ConnectionError("Transport disconnected")
        # Deep-copy through JSON round-trip for realism
        self.sent.append(json.loads(json.dumps(msg)))

    def make_disconnected(self):
        """Helper to simulate disconnection."""
        self._connected = False


class MockRedis:
    """Mock Redis client for testing Redis backend."""

    def __init__(self):
        self.data = {}
        self.published = []
        self.should_fail = False

    @classmethod
    def from_url(cls, url):
        return cls()

    async def setex(self, key, ttl, value):
        if self.should_fail:
            raise Exception("Redis connection failed")
        self.data[key] = {"value": value, "ttl": ttl}

    async def exists(self, key):
        if self.should_fail:
            raise Exception("Redis connection failed")
        return key in self.data

    async def delete(self, key):
        if self.should_fail:
            raise Exception("Redis connection failed")
        self.data.pop(key, None)

    async def expire(self, key, ttl):
        if self.should_fail:
            raise Exception("Redis connection failed")
        if key in self.data:
            self.data[key]["ttl"] = ttl

    async def publish(self, channel, message):
        if self.should_fail:
            raise Exception("Redis connection failed")
        self.published.append({"channel": channel, "message": message})

    def pubsub(self):
        return MockPubSub()

    def close(self):
        pass


class MockPubSub:
    """Mock Redis PubSub."""

    def __init__(self):
        self.subscribed_channels = set()

    async def subscribe(self, channel):
        self.subscribed_channels.add(channel)

    async def unsubscribe(self, channel):
        self.subscribed_channels.discard(channel)

    async def listen(self):
        # Simulate empty message stream
        if False:  # Never yield anything
            yield {}

    async def aclose(self):
        """Async close for redis.asyncio compatibility."""
        pass

    def close(self):
        pass


# --------------------------------------------------------------------------- #
# SessionRegistry fixture (memory backend)                                    #
# --------------------------------------------------------------------------- #
@pytest.fixture()
async def registry() -> SessionRegistry:
    """Initialise an in-memory SessionRegistry and tear it down after tests."""
    reg = SessionRegistry(backend="memory")
    await reg.initialize()
    yield reg
    await reg.shutdown()


# --------------------------------------------------------------------------- #
# Core CRUD behaviour                                                         #
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_add_get_remove(registry: SessionRegistry):
    """Add ➜ get (async & sync) ➜ remove and verify cache/state."""
    tr = FakeSSETransport("A")
    await registry.add_session("A", tr)

    assert await registry.get_session("A") is tr
    assert registry.get_session_sync("A") is tr
    assert await registry.get_session("missing") is None

    # Remove twice - second call must be harmless
    await registry.remove_session("A")
    await registry.remove_session("A")

    assert not await tr.is_connected()
    assert registry.get_session_sync("A") is None


# --------------------------------------------------------------------------- #
# broadcast ➜ respond with different payload types                            #
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "payload",
    [
        {"method": "ping", "id": 1, "params": {}},  # dict
        ["x", "y", 42],  # list
        "plain-string",  # str
    ],
)
async def test_broadcast_and_respond(payload, monkeypatch, registry: SessionRegistry):
    """broadcast stores the payload; respond routes it to generate_response."""
    tr = FakeSSETransport("B")
    await registry.add_session("B", tr)

    captured: Dict[str, Any] = {}

    # Patch generate_response so we can verify it's called with our payload:
    async def fake_generate_response(*, message, transport, **kwargs):  # noqa: D401
        captured["message"] = message
        captured["transport"] = transport
        captured["kwargs"] = kwargs

    monkeypatch.setattr(registry, "generate_response", fake_generate_response)

    await registry.broadcast("B", payload)
    await registry.respond(server_id=None, user={}, session_id="B")

    assert captured["transport"] is tr
    assert captured["message"] == payload


async def test_broadcast_redis_input(monkeypatch, registry: SessionRegistry):
    """test input to publish for redis"""

    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)
    registry._backend = "redis"

    mock_redis = AsyncMock()
    registry._redis = mock_redis

    fixed_ts = 1_234_567.890
    monkeypatch.setattr("mcpgateway.cache.session_registry.time.time", lambda: fixed_ts)

    msg = {"a": 1}
    expected_payload = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts})

    await registry.broadcast("B", msg)

    mock_redis.publish.assert_awaited_once_with("B", expected_payload)

    mock_redis.publish.reset_mock()

    msg = ["a", "b", "c"]
    expected_payload = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts})

    await registry.broadcast("B", msg)

    mock_redis.publish.assert_awaited_once_with("B", expected_payload)

    mock_redis.publish.reset_mock()

    msg = 123
    expected_payload = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts})

    await registry.broadcast("B", msg)

    mock_redis.publish.assert_awaited_once_with("B", expected_payload)

    mock_redis.publish.reset_mock()

    msg = "hello\nworld"
    expected_payload = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts})

    await registry.broadcast("B", msg)

    mock_redis.publish.assert_awaited_once_with("B", expected_payload)

    mock_redis.publish.reset_mock()


async def test_broadcast_database_input(monkeypatch, registry: SessionRegistry, caplog):
    """test input to publish for database"""

    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)
    registry._backend = "database"

    mock_db = MagicMock()  # Use MagicMock for sync SQLAlchemy session methods
    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)
    monkeypatch.setattr("mcpgateway.cache.session_registry.get_db", lambda: iter([mock_db]), raising=True)

    monkeypatch.setattr("asyncio.to_thread", lambda func, *a, **k: func(*a, **k))

    fixed_ts = 1_234_567.890
    monkeypatch.setattr("mcpgateway.cache.session_registry.time.time", lambda: fixed_ts)

    msg = {"a": 1}
    expected_msg_json = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts}).decode()

    await registry.broadcast("B", msg)

    assert mock_db.add.call_count == 1
    actual_record = mock_db.add.call_args[0][0]
    assert isinstance(actual_record, SessionMessageRecord)
    assert actual_record.session_id == "B"
    assert actual_record.message == expected_msg_json

    mock_db.add.reset_mock()

    msg = ["a", "b", "c"]
    expected_msg_json = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts}).decode()

    await registry.broadcast("B", msg)

    assert mock_db.add.call_count == 1
    actual_record = mock_db.add.call_args[0][0]
    assert isinstance(actual_record, SessionMessageRecord)
    assert actual_record.session_id == "B"
    assert actual_record.message == expected_msg_json

    mock_db.add.reset_mock()

    msg = 123
    expected_msg_json = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts}).decode()

    await registry.broadcast("B", msg)

    assert mock_db.add.call_count == 1
    actual_record = mock_db.add.call_args[0][0]
    assert isinstance(actual_record, SessionMessageRecord)
    assert actual_record.session_id == "B"
    assert actual_record.message == expected_msg_json

    mock_db.add.reset_mock()

    msg = "hello\nworld"
    expected_msg_json = orjson.dumps({"type": "message", "message": msg, "timestamp": fixed_ts}).decode()

    await registry.broadcast("B", msg)

    assert mock_db.add.call_count == 1
    actual_record = mock_db.add.call_args[0][0]
    assert isinstance(actual_record, SessionMessageRecord)
    assert actual_record.session_id == "B"
    assert actual_record.message == expected_msg_json

    mock_db.add.reset_mock()
    mock_db.commit = Mock(side_effect=Exception("db error"))

    msg = "hello\nworld"
    expected_msg_json = json.dumps({"type": "message", "message": msg, "timestamp": fixed_ts})

    await registry.broadcast("B", msg)

    mock_db.rollback.assert_called_once()

    assert "Database error during broadcast" in caplog.text


# --------------------------------------------------------------------------- #
# Fixtures to stub get_db and the three *Service objects                      #
# --------------------------------------------------------------------------- #
@pytest.fixture()
def stub_db(monkeypatch):
    """Patch ``get_db`` to return a synchronous dummy iterator."""

    def _dummy_iter():
        yield None

    monkeypatch.setattr(
        "mcpgateway.cache.session_registry.get_db",
        lambda: _dummy_iter(),
        raising=False,
    )


@pytest.fixture()
def stub_services(monkeypatch):
    """Replace list_* service methods so they return predictable data."""

    class _Item:
        def model_dump(self, *_, **__) -> Dict[str, str]:  # noqa: D401
            return {"name": "demo"}

    async def _return_items(*args, **kwargs):  # noqa: D401
        return [_Item()]

    mod = "mcpgateway.cache.session_registry"
    monkeypatch.setattr(f"{mod}.tool_service.list_tools", _return_items, raising=False)
    monkeypatch.setattr(f"{mod}.tool_service.list_server_tools", _return_items, raising=False)
    monkeypatch.setattr(f"{mod}.prompt_service.list_prompts", _return_items, raising=False)
    monkeypatch.setattr(f"{mod}.prompt_service.list_server_prompts", _return_items, raising=False)
    monkeypatch.setattr(f"{mod}.resource_service.list_resources", _return_items, raising=False)
    monkeypatch.setattr(f"{mod}.resource_service.list_server_resources", _return_items, raising=False)


def test_redis_importerror_isolated():
    # Backup original sys.modules state
    original_redis_asyncio = sys.modules.get("redis.asyncio")
    original_my_module = sys.modules.get("mcpgateway.cache.session_registry")

    # Simulate ImportError for redis.asyncio
    with patch.dict(sys.modules, {"redis.asyncio": None}):
        # if 'mcpgateway.cache.session_registry' in sys.modules:
        #     del sys.modules['mcpgateway.cache.session_registry']  # Force re-import

        # First-Party
        import mcpgateway.cache.session_registry

        importlib.reload(mcpgateway.cache.session_registry)
        assert not mcpgateway.cache.session_registry.REDIS_AVAILABLE

    # Cleanup: restore the original sys.modules entries
    if original_redis_asyncio is not None:
        sys.modules["redis.asyncio"] = original_redis_asyncio
    else:
        sys.modules.pop("redis.asyncio", None)

    if original_my_module is not None:
        sys.modules["mcpgateway.cache.session_registry"] = original_my_module
    else:
        sys.modules.pop("mcpgateway.cache.session_registry", None)


def test_sqlalchemy_importerror_isolated():
    # Backup original sys.modules state
    original_sqlalchemy = sys.modules.get("sqlalchemy")
    original_my_module = sys.modules.get("mcpgateway.cache.session_registry")

    # Simulate ImportError for redis.asyncio
    with patch.dict(sys.modules, {"sqlalchemy": None}):
        # if 'mcpgateway.cache.session_registry' in sys.modules:
        # del sys.modules['mcpgateway.cache.session_registry']  # Force re-import

        # First-Party
        import mcpgateway.cache.session_registry

        importlib.reload(mcpgateway.cache.session_registry)
        assert not mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE

    # Cleanup: restore the original sys.modules entries and reload to reset SQLALCHEMY_AVAILABLE
    if original_sqlalchemy is not None:
        sys.modules["sqlalchemy"] = original_sqlalchemy
    else:
        sys.modules.pop("sqlalchemy", None)

    if original_my_module is not None:
        sys.modules["mcpgateway.cache.session_registry"] = original_my_module
    else:
        sys.modules.pop("mcpgateway.cache.session_registry", None)

    import mcpgateway.cache.session_registry

    importlib.reload(mcpgateway.cache.session_registry)


# --------------------------------------------------------------------------- #
# generate_response branches                                                  #
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_generate_response_initialize(registry: SessionRegistry):
    """The *initialize* branch sends result + notifications (>= 5 messages)."""
    tr = FakeSSETransport("init")
    await registry.add_session("init", tr)

    msg = {"method": "initialize", "id": 101, "params": {"protocol_version": settings.protocol_version}}

    mock_response = Mock()
    mock_response.json.return_value = {"result": {"protocolVersion": settings.protocol_version}, "id": 101}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id=None,
            user={"token": "test"},
        )

    # Implementation may emit 5 or 6 messages (roots/list_changed optional)
    assert len(tr.sent) >= 5

    first = tr.sent[0]
    assert first["id"] == 101
    assert first["result"]["protocolVersion"] == settings.protocol_version
    assert re.match(r"notifications/initialized$", tr.sent[1]["method"])


@pytest.mark.asyncio
async def test_generate_response_ping(registry: SessionRegistry):
    """The *ping* branch should echo an empty result."""
    tr = FakeSSETransport("ping")
    await registry.add_session("ping", tr)

    msg = {"method": "ping", "id": 77, "params": {}}

    mock_response = Mock()
    mock_response.json.return_value = {"result": {}, "id": 77}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id=None,
            user={"token": "test"},
        )

    assert tr.sent[-1] == {"jsonrpc": "2.0", "result": {}, "id": 77}


@pytest.mark.asyncio
async def test_generate_response_tools_list(registry: SessionRegistry, stub_db, stub_services):
    """*tools/list* responds with the stubbed ToolService payload."""
    tr = FakeSSETransport("tools")
    await registry.add_session("tools", tr)

    msg = {"method": "tools/list", "id": 42, "params": {}}

    mock_response = Mock()
    mock_response.json.return_value = {"jsonrpc": "2.0", "result": [{"name": "demo"}], "id": 42}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id=None,
            user={"token": "test"},
        )

    reply = tr.sent[-1]
    print(f"{reply=}")
    assert reply["id"] == 42
    assert reply["result"] == [{"name": "demo"}]


@pytest.mark.asyncio
async def test_generate_response_resources_list(registry: SessionRegistry, stub_db, stub_services):
    """*resources/list* responds with the stubbed ResourceService payload."""
    tr = FakeSSETransport("resources")
    await registry.add_session("resources", tr)

    msg = {"method": "resources/list", "id": 43, "params": {}}

    mock_response = Mock()
    mock_response.json.return_value = {"jsonrpc": "2.0", "result": [{"name": "demo"}], "id": 42}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id=None,
            user={"token": "test"},
        )

    reply = tr.sent[-1]
    assert reply["id"] == 43
    assert reply["result"] == [{"name": "demo"}]


@pytest.mark.asyncio
async def test_generate_response_prompts_list(registry: SessionRegistry, stub_db, stub_services):
    """*prompts/list* responds with the stubbed PromptService payload."""
    tr = FakeSSETransport("prompts")
    await registry.add_session("prompts", tr)

    msg = {"method": "prompts/list", "id": 44, "params": {}}

    mock_response = Mock()
    mock_response.json.return_value = {"jsonrpc": "2.0", "result": [{"name": "demo"}], "id": 42}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id=None,
            user={"token": "test"},
        )

    reply = tr.sent[-1]
    assert reply["id"] == 44
    assert reply["result"] == [{"name": "demo"}]


@pytest.mark.asyncio
async def test_generate_response_tools_call(registry: SessionRegistry, stub_db, stub_services):
    """*tools/call* makes HTTP request and returns response."""
    tr = FakeSSETransport("tools_call")
    await registry.add_session("tools_call", tr)

    # Mock httpx.AsyncClient properly as an async context manager
    mock_response = Mock()
    mock_response.json.return_value = {"result": "tool_executed"}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    # Create a proper async context manager mock
    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        msg = {"method": "tools/call", "id": 45, "params": {"name": "test_tool", "arguments": {"arg1": "value1"}}}
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id=None,
            user={"token": "test_token"},
        )

    reply = tr.sent[-1]
    assert reply["id"] == 45
    assert reply["result"] == "tool_executed"


@pytest.mark.asyncio
async def test_generate_response_server_specific_tools_list(registry: SessionRegistry, stub_db, stub_services):
    """*tools/list* with server_id calls server-specific method."""
    tr = FakeSSETransport("server_tools")
    await registry.add_session("server_tools", tr)

    msg = {"method": "tools/list", "id": 46, "params": {}}

    mock_response = Mock()
    mock_response.json.return_value = {"jsonrpc": "2.0", "result": [{"name": "demo"}], "id": 46}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id="server123",
            user={"token": "test"},
        )

    reply = tr.sent[-1]
    assert reply["id"] == 46
    assert reply["result"] == [{"name": "demo"}]


@pytest.mark.asyncio
async def test_generate_response_server_specific_resources_list(registry: SessionRegistry, stub_db, stub_services):
    """*resources/list* responds with server_id calls server-specific method."""
    tr = FakeSSETransport("resources")
    await registry.add_session("resources", tr)

    msg = {"method": "resources/list", "id": 43, "params": {}}

    mock_response = Mock()
    mock_response.json.return_value = {"jsonrpc": "2.0", "result": [{"name": "demo"}], "id": 43}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id="server123",
            user={"token": "test"},
        )

    reply = tr.sent[-1]
    assert reply["id"] == 43
    assert reply["result"] == [{"name": "demo"}]


@pytest.mark.asyncio
async def test_generate_response_server_specific_prompts_list(registry: SessionRegistry, stub_db, stub_services):
    """*prompts/list* responds with server_id calls server-specific method."""
    tr = FakeSSETransport("prompts")
    await registry.add_session("prompts", tr)

    msg = {"method": "prompts/list", "id": 44, "params": {}}

    mock_response = Mock()
    mock_response.json.return_value = {"jsonrpc": "2.0", "result": [{"name": "demo"}], "id": 44}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id="server123",
            user={"token": "test"},
        )

    reply = tr.sent[-1]
    assert reply["id"] == 44
    assert reply["result"] == [{"name": "demo"}]


@pytest.mark.asyncio
async def test_generate_response_unknown_method(registry: SessionRegistry, stub_db):
    """Unknown method returns empty result."""
    tr = FakeSSETransport("unknown")
    await registry.add_session("unknown", tr)

    mock_response = Mock()
    mock_response.json.return_value = {"result": {}}

    mock_client = AsyncMock()
    mock_client.post.return_value = mock_response

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return mock_client

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    msg = {"method": "unknown_method", "id": 47, "params": {}}
    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(
            message=msg,
            transport=tr,
            server_id=None,
            user={"token": "test"},
        )

    reply = tr.sent[-1]
    print(f"{reply=}")
    assert reply["id"] == 47
    assert reply["result"] == {}


@pytest.mark.asyncio
async def test_generate_response_no_method_or_id(registry: SessionRegistry):
    """Message without method or id is ignored."""
    tr = FakeSSETransport("no_method")
    await registry.add_session("no_method", tr)

    msg = {"some": "data"}
    await registry.generate_response(
        message=msg,
        transport=tr,
        server_id=None,
        user={},
    )

    # Should not send any response
    assert len(tr.sent) == 0


# --------------------------------------------------------------------------- #
# handle_initialize_logic success & errors                                    #
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_handle_initialize_success(registry: SessionRegistry):
    body = {"protocol_version": settings.protocol_version}
    res = await registry.handle_initialize_logic(body)
    assert res.protocol_version == settings.protocol_version


@pytest.mark.asyncio
async def test_handle_initialize_missing_version_error(registry: SessionRegistry):
    with pytest.raises(HTTPException) as exc:
        await registry.handle_initialize_logic({})
    assert exc.value.headers["MCP-Error-Code"] == "-32002"


@pytest.mark.asyncio
async def test_handle_initialize_unsupported_version_warning(registry: SessionRegistry, caplog):
    caplog.set_level(logging.WARNING, logger="mcpgateway.cache.session_registry")
    body = {"protocol_version": "999"}

    await registry.handle_initialize_logic(body)

    assert "Using non default protocol version: 999" in caplog.text


@pytest.mark.asyncio
async def test_handle_initialize_creates_observability_span(registry: SessionRegistry, monkeypatch):
    entered = []

    class _SpanCtx:
        def __enter__(self):
            entered.append("enter")
            return self

        def __exit__(self, exc_type, exc, tb):
            entered.append("exit")
            return False

    create_span = MagicMock()

    def _fake_create_span(name, attributes):
        create_span(name, attributes)
        return _SpanCtx()

    monkeypatch.setattr("mcpgateway.observability.create_span", _fake_create_span)

    body = {"protocol_version": settings.protocol_version}
    res = await registry.handle_initialize_logic(body, session_id="sess-1", server_id="server-1")

    assert res.protocol_version == settings.protocol_version
    assert entered == ["enter", "exit"]
    create_span.assert_called_once_with(
        "mcp.initialize",
        {
            "mcp.protocol_version": settings.protocol_version,
            "mcp.session_id": "sess-1",
            "server.id": "server-1",
        },
    )


# --------------------------------------------------------------------------- #
# Backend initialization tests                                                #
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_none_backend():
    """Test 'none' backend does no-op for all operations."""
    registry = SessionRegistry(backend="none")
    await registry.initialize()

    try:
        tr = FakeSSETransport("none_test")

        # All operations should be no-ops
        await registry.add_session("none_test", tr)
        assert registry.get_session_sync("none_test") is None
        assert await registry.get_session("none_test") is None

        await registry.remove_session("none_test")
        assert not tr.disconnect_called

        await registry.broadcast("none_test", {"test": "message"})
        await registry.respond(server_id=None, user={"token": "test"}, session_id="none_test")

        assert len(tr.sent) == 0
    finally:
        await registry.shutdown()


@pytest.mark.asyncio
async def test_redis_backend_init_no_redis_available(monkeypatch):
    """Test Redis backend when Redis not available."""
    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", False)

    with pytest.raises(ValueError, match="Redis backend requested but redis package not installed"):
        SessionRegistry(backend="redis", redis_url="redis://localhost:6379")


@pytest.mark.asyncio
async def test_redis_backend_init_no_url(monkeypatch):
    """Test Redis backend without URL."""
    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    with pytest.raises(ValueError, match="Redis backend requires redis_url"):
        SessionRegistry(backend="redis")


@pytest.mark.asyncio
async def test_database_backend_init_no_sqlalchemy_available(monkeypatch):
    """Test database backend when SQLAlchemy not available."""
    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", False)

    with pytest.raises(ValueError, match="Database backend requested but SQLAlchemy not installed"):
        SessionRegistry(backend="database", database_url="sqlite:///test.db")


@pytest.mark.asyncio
async def test_database_backend_init_no_url(monkeypatch):
    """Test database backend without URL."""
    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)

    with pytest.raises(ValueError, match="Database backend requires database_url"):
        SessionRegistry(backend="database")


@pytest.mark.asyncio
async def test_invalid_backend():
    """Test initialization with invalid backend."""
    with pytest.raises(ValueError, match="Invalid backend: invalid"):
        SessionRegistry(backend="invalid")


# --------------------------------------------------------------------------- #
# Redis backend session operations                                            #
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_redis_session_operations(monkeypatch):
    """Test Redis backend session operations."""
    mock_redis = MockRedis()

    # Patch Redis imports before creating the registry
    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    # Patch the get_redis_client function used by session_registry
    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(backend="redis", redis_url="redis://localhost:6379")
        await registry.initialize()

        try:
            tr = FakeSSETransport("redis_session")

            # Add session
            await registry.add_session("redis_session", tr)
            assert registry.get_session_sync("redis_session") is tr
            assert "mcp:session:redis_session" in mock_redis.data
            assert len(mock_redis.published) == 1

            # Get session from local cache
            result = await registry.get_session("redis_session")
            assert result is tr

            # Remove session
            await registry.remove_session("redis_session")
            assert registry.get_session_sync("redis_session") is None
            assert tr.disconnect_called

            # Test broadcast
            message = {"method": "test", "params": {}}
            await registry.broadcast("test_session", message)
            assert len(mock_redis.published) >= 2

        finally:
            await registry.shutdown()


@pytest.mark.asyncio
async def test_redis_error_handling(monkeypatch, caplog):
    """Test Redis backend error handling."""
    mock_redis = MockRedis()
    mock_redis.should_fail = True

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    # Patch the get_redis_client function used by session_registry
    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(backend="redis", redis_url="redis://localhost:6379")
        await registry.initialize()

        try:
            tr = FakeSSETransport("redis_error")

            # Operations should not raise exceptions despite Redis failures
            await registry.add_session("redis_error", tr)
            await registry.get_session("redis_error")
            await registry.remove_session("redis_error")
            await registry.broadcast("redis_error", {"test": "message"})

        finally:
            await registry.shutdown()

        assert "Redis error adding session redis_error" in caplog.text


# --------------------------------------------------------------------------- #
# Database backend session operations                                         #
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_database_session_operations(monkeypatch):
    """Test database backend session operations."""
    mock_db_session = Mock()
    mock_db_session.query.return_value.filter.return_value.first.return_value = None
    mock_db_session.query.return_value.filter.return_value.delete.return_value = 0

    async def immediate_execution(func):
        return func()

    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)
    monkeypatch.setattr("mcpgateway.cache.session_registry.get_db", lambda: iter([mock_db_session]))
    monkeypatch.setattr("mcpgateway.cache.session_registry.asyncio.to_thread", immediate_execution)

    registry = SessionRegistry(backend="database", database_url="sqlite:///test.db")
    await registry.initialize()

    try:
        tr = FakeSSETransport("db_session")

        # Add session
        await registry.add_session("db_session", tr)
        assert registry.get_session_sync("db_session") is tr

        # Get session (should return None since not in DB)
        result = await registry.get_session("db_session")
        assert result is tr  # Found in local cache

        # Remove session
        await registry.remove_session("db_session")
        assert registry.get_session_sync("db_session") is None
        assert tr.disconnect_called

        # Test broadcast
        message = {"method": "test", "params": {}}
        await registry.broadcast("test_session", message)

    finally:
        await registry.shutdown()


@pytest.mark.asyncio
async def test_database_add_session_exception(monkeypatch, caplog):
    """Test database backend session operations."""
    mock_db_session = Mock()
    mock_db_session.add = Mock()  # okay
    mock_db_session.commit = Mock()
    mock_db_session.rollback = Mock()
    mock_db_session.close = Mock()

    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)
    monkeypatch.setattr("mcpgateway.cache.session_registry.get_db", lambda: iter([mock_db_session]), raising=True)

    monkeypatch.setattr("asyncio.to_thread", lambda func, *a, **k: func(*a, **k))

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    registry = SessionRegistry(backend="database", database_url="sqlite:///test.db")
    await registry.initialize()

    mock_db_session.commit = Mock(side_effect=Exception("db error"))
    mock_db_session.rollback.reset_mock()
    mock_db_session.close.reset_mock()

    tr = FakeSSETransport("db_session")
    await registry.add_session("db_session", tr)

    mock_db_session.rollback.assert_called_once()
    mock_db_session.close.assert_called_once()

    assert "Database error adding session db_session" in caplog.text


@pytest.mark.asyncio
async def test_database_remove_session_exception(monkeypatch, caplog):
    """Test database backend session operations."""
    mock_db_session = Mock()
    mock_db_session.filter.return_value.delete = Mock()
    mock_db_session.query = Mock(return_value=mock_db_session)
    mock_db_session.commit = Mock(side_effect=Exception("db error"))
    mock_db_session.rollback = Mock()
    mock_db_session.close = Mock()

    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)
    monkeypatch.setattr("mcpgateway.cache.session_registry.get_db", lambda: iter([mock_db_session]), raising=True)

    monkeypatch.setattr("asyncio.to_thread", lambda func, *a, **k: func(*a, **k))

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    registry = SessionRegistry(backend="database", database_url="sqlite:///test.db")
    await registry.initialize()

    mock_db_session.rollback.reset_mock()
    mock_db_session.close.reset_mock()

    tr = FakeSSETransport("db_session")
    tr.disconnect = AsyncMock()

    monkeypatch.setattr(registry, "_sessions", {"db_session": tr})

    await registry.remove_session("db_session")

    tr.disconnect.assert_awaited_once()

    # 8) And the DB path hit the commit‐>exception branch:
    mock_db_session.rollback.assert_called_once()
    mock_db_session.close.assert_called_once()

    assert "Database error removing session db_session" in caplog.text


# --------------------------------------------------------------------------- #
# Cleanup and error scenarios                                                 #
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_transport_disconnect_error(registry: SessionRegistry):
    """Test handling transport disconnect errors."""
    tr = FakeSSETransport("error_test")

    # Mock disconnect to raise exception
    async def failing_disconnect():
        raise Exception("Disconnect failed")

    tr.disconnect = failing_disconnect

    await registry.add_session("error_test", tr)

    # Should not raise exception
    await registry.remove_session("error_test")


@pytest.mark.asyncio
async def test_concurrent_session_operations():
    """Test concurrent session operations."""
    registry = SessionRegistry(backend="memory")
    await registry.initialize()

    try:

        async def add_remove_session(session_id):
            transport = FakeSSETransport(session_id)
            await registry.add_session(session_id, transport)
            await asyncio.sleep(0.001)  # Small delay
            await registry.remove_session(session_id)

        # Run multiple concurrent operations
        tasks = [add_remove_session(f"session_{i}") for i in range(5)]
        await asyncio.gather(*tasks)

        # All sessions should be cleaned up
        assert len(registry._sessions) == 0

    finally:
        await registry.shutdown()


@pytest.mark.asyncio
async def test_memory_cleanup_task():
    """Test memory cleanup task removes disconnected sessions."""
    registry = SessionRegistry(backend="memory")
    await registry.initialize()

    try:
        tr = FakeSSETransport("cleanup_test")
        await registry.add_session("cleanup_test", tr)

        # Simulate disconnection
        tr.make_disconnected()

        # Manually trigger cleanup logic
        async with registry._lock:
            local_transports = registry._sessions.copy()

        for session_id, transport in local_transports.items():
            if not await transport.is_connected():
                await registry.remove_session(session_id)

        assert registry.get_session_sync("cleanup_test") is None

    finally:
        await registry.shutdown()


@pytest.mark.asyncio
async def test_redis_shutdown(monkeypatch):
    """shutdown() should close PubSub but not the shared Redis client."""

    # Tell the registry that the Redis extras are present
    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # ── fake PubSub object ────────────────────────────────────────────────
    mock_pubsub = AsyncMock(name="MockPubSub")
    mock_pubsub.aclose = AsyncMock()

    # ── fake Redis client ────────────────────────────────────────────────
    mock_redis = AsyncMock(name="MockRedis")
    # pubsub() is **not** awaited in prod code, so a plain Mock is fine
    mock_redis.pubsub = Mock(return_value=mock_pubsub)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    # ── patch the get_redis_client function ────────────────────────
    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(
            backend="redis",
            redis_url="redis://localhost:6379",
        )
        await registry.initialize()  # calls mock_redis.pubsub()

        await registry.shutdown()

        # Only PubSub should be closed, not the shared Redis client
        mock_pubsub.aclose.assert_awaited_once()


@pytest.mark.asyncio
async def test_shutdown_with_redis_error(monkeypatch):
    """shutdown() should swallow PubSub aclose() errors."""

    # Tell the registry that the Redis extras are present
    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # ── fake PubSub object ────────────────────────────────────────────────
    mock_pubsub = AsyncMock(name="MockPubSub")
    mock_pubsub.aclose = AsyncMock(side_effect=Exception("PubSub close error"))

    # ── fake Redis client ────────────────────────────────────────────────
    mock_redis = AsyncMock(name="MockRedis")
    # pubsub() is **not** awaited in prod code, so a plain Mock is fine
    mock_redis.pubsub = Mock(return_value=mock_pubsub)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    # ── patch the get_redis_client function ────────────────────────
    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(
            backend="redis",
            redis_url="redis://localhost:6379",
        )
        await registry.initialize()  # calls mock_redis.pubsub()

        # must swallow PubSub aclose() exception
        await registry.shutdown()


@pytest.mark.asyncio
async def test_full_memory_workflow(stub_db, stub_services):
    """Test complete workflow with memory backend."""
    registry = SessionRegistry(backend="memory")
    await registry.initialize()

    try:
        # Add session
        transport = FakeSSETransport("workflow_test")
        await registry.add_session("workflow_test", transport)

        # Broadcast message
        init_message = {"method": "initialize", "id": 1, "params": {"protocol_version": settings.protocol_version}}
        await registry.broadcast("workflow_test", init_message)

        mock_response = Mock()
        mock_response.json.return_value = {"result": {"protocolVersion": settings.protocol_version}, "id": 1}

        mock_client = AsyncMock()
        mock_client.post.return_value = mock_response

        class MockAsyncClient:
            def __init__(self, *args, **kwargs):
                pass

            async def __aenter__(self):
                return mock_client

            async def __aexit__(self, exc_type, exc_val, exc_tb):
                return None

        # Respond to message
        with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
            await registry.respond(server_id=None, user={"token": "test"}, session_id="workflow_test")

        # Should have received initialize response + notifications
        assert len(transport.sent) >= 5

        # Verify initialize response
        init_response = transport.sent[0]
        assert init_response["id"] == 1
        assert "protocolVersion" in init_response["result"]

        # Clean up
        await registry.remove_session("workflow_test")
        assert transport.disconnect_called

    finally:
        await registry.shutdown()


@pytest.mark.asyncio
async def test_database_get_session_exception(monkeypatch, caplog):
    """Test database backend get_session with exception."""
    mock_db_session = Mock()
    mock_db_session.query.return_value.filter.return_value.first = Mock(side_effect=Exception("db error"))
    mock_db_session.rollback = Mock()
    mock_db_session.close = Mock()

    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)
    monkeypatch.setattr("mcpgateway.cache.session_registry.get_db", lambda: iter([mock_db_session]), raising=True)
    monkeypatch.setattr("asyncio.to_thread", lambda func, *a, **k: func(*a, **k))

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    registry = SessionRegistry(backend="database", database_url="sqlite:///test.db")
    await registry.initialize()

    result = await registry.get_session("test_session")
    assert result is None
    assert "Database error checking session test_session" in caplog.text


@pytest.mark.asyncio
async def test_database_get_session_exists_in_db(monkeypatch, caplog):
    """Test database backend get_session when session exists in DB but not locally."""
    mock_record = Mock()
    mock_db_session = Mock()
    mock_db_session.query.return_value.filter.return_value.first.return_value = mock_record
    mock_db_session.close = Mock()

    monkeypatch.setattr("mcpgateway.cache.session_registry.SQLALCHEMY_AVAILABLE", True)
    monkeypatch.setattr("mcpgateway.cache.session_registry.get_db", lambda: iter([mock_db_session]), raising=True)

    # Mock asyncio.to_thread to return the result directly
    async def mock_to_thread(func, *args, **kwargs):
        return func(*args, **kwargs)

    monkeypatch.setattr("asyncio.to_thread", mock_to_thread)

    registry = SessionRegistry(backend="database", database_url="sqlite:///test.db")
    await registry.initialize()

    caplog.set_level(logging.INFO, logger="mcpgateway.cache.session_registry")
    result = await registry.get_session("test_session")
    assert result is None
    assert "Session test_session exists in database but not in local cache" in caplog.text


@pytest.mark.asyncio
async def test_redis_get_session_exists_in_redis(monkeypatch, caplog):
    """Test Redis backend get_session when session exists in Redis but not locally."""
    mock_redis = MockRedis()
    mock_redis.data["mcp:session:test_session"] = {"value": "1", "ttl": 3600}

    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    caplog.set_level(logging.INFO, logger="mcpgateway.cache.session_registry")

    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(backend="redis", redis_url="redis://localhost:6379")
        await registry.initialize()

        result = await registry.get_session("test_session")
        assert result is None
        assert "Session test_session exists in Redis but not in local cache" in caplog.text

        await registry.shutdown()


@pytest.mark.asyncio
async def test_redis_get_session_exception(monkeypatch, caplog):
    """Test Redis backend get_session with exception."""
    mock_redis = MockRedis()
    mock_redis.should_fail = True

    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(backend="redis", redis_url="redis://localhost:6379")
        await registry.initialize()

        result = await registry.get_session("test_session")
        assert result is None
        assert "Redis error checking session test_session" in caplog.text

        await registry.shutdown()


@pytest.mark.asyncio
async def test_remove_session_transport_disconnect_error(registry: SessionRegistry, caplog):
    """Test remove_session when transport.disconnect() raises exception."""
    tr = FakeSSETransport("error_session")
    tr.disconnect = AsyncMock(side_effect=Exception("disconnect error"))

    await registry.add_session("error_session", tr)

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    await registry.remove_session("error_session")

    assert "Error disconnecting transport for session error_session" in caplog.text
    assert registry.get_session_sync("error_session") is None


@pytest.mark.asyncio
async def test_respond_memory_backend_no_message(registry: SessionRegistry):
    """Test respond with memory backend when no message is stored."""
    tr = FakeSSETransport("test_session")
    await registry.add_session("test_session", tr)

    # Ensure no message is stored
    registry._session_message = None

    # The respond method should handle None _session_message gracefully
    await registry.respond(server_id=None, user={"token": "test"}, session_id="test_session")


@pytest.mark.asyncio
async def test_respond_memory_backend_with_none_message_content(registry: SessionRegistry, caplog):
    """Test respond with memory backend when message content is None (not the whole dict)."""
    caplog.set_level(logging.WARNING, logger="mcpgateway.cache.session_registry")

    tr = FakeSSETransport("test_session")
    await registry.add_session("test_session", tr)

    # Set _session_message to a dict with "message" key having None value
    # This is the specific edge case that was previously causing orjson.loads(None) to crash
    registry._session_message = {"session_id": "test_session", "message": None}

    with patch.object(registry, "generate_response", new_callable=AsyncMock) as mock_gen:
        await registry.respond(server_id=None, user={"token": "test"}, session_id="test_session")

        # Should NOT call generate_response since message content is None
        mock_gen.assert_not_called()

    # Should log a warning about the None message content
    assert "message content is None" in caplog.text


@pytest.mark.asyncio
async def test_respond_memory_backend_with_message_no_transport(registry: SessionRegistry):
    """Test respond with memory backend when message exists but no transport."""
    registry._session_message = {"session_id": "missing_session", "message": json.dumps({"method": "ping", "id": 1})}

    with patch.object(registry, "generate_response", new_callable=AsyncMock) as mock_gen:
        await registry.respond(server_id=None, user={"token": "test"}, session_id="missing_session")

        mock_gen.assert_not_called()


@pytest.mark.asyncio
async def test_respond_memory_backend_with_session_message_check(registry: SessionRegistry):
    """Test respond memory backend checks _session_message properly."""
    tr = FakeSSETransport("test_session")
    await registry.add_session("test_session", tr)

    # Set up a message but without transport
    registry._session_message = {"session_id": "test_session", "message": json.dumps({"method": "ping", "id": 1})}

    with patch.object(registry, "generate_response", new_callable=AsyncMock) as mock_gen:
        await registry.respond(server_id=None, user={"token": "test"}, session_id="test_session")

        # Should call generate_response since transport exists
        mock_gen.assert_called_once()
        args, kwargs = mock_gen.call_args
        assert kwargs["message"] == {"method": "ping", "id": 1}
        assert kwargs["transport"] is tr


@pytest.mark.asyncio
async def test_generate_response_jsonrpc_error(registry: SessionRegistry):
    """Test generate_response with JSONRPCError."""
    tr = FakeSSETransport("jsonrpc_error")
    await registry.add_session("jsonrpc_error", tr)

    message = {"method": "test_method", "id": 1, "params": {}}

    # Mock ResilientHttpClient to raise JSONRPCError
    # First-Party
    from mcpgateway.validation.jsonrpc import JSONRPCError

    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            raise JSONRPCError(code=-32601, message="Method not found")

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(message=message, transport=tr, server_id=None, user={"token": "test"})

    # Should have sent error response
    assert len(tr.sent) == 1
    response = tr.sent[0]
    assert "error" in response
    assert response["id"] == 1


@pytest.mark.asyncio
async def test_generate_response_general_exception(registry: SessionRegistry):
    """Test generate_response with general exception."""
    tr = FakeSSETransport("general_error")
    await registry.add_session("general_error", tr)

    message = {"method": "test_method", "id": 1, "params": {}}

    # Mock ResilientHttpClient to raise general exception
    class MockAsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            raise Exception("Network error")

        async def __aexit__(self, exc_type, exc_val, exc_tb):
            return None

    with patch("mcpgateway.cache.session_registry.ResilientHttpClient", MockAsyncClient):
        await registry.generate_response(message=message, transport=tr, server_id=None, user={"token": "test"})

    # Should have sent error response
    assert len(tr.sent) == 1
    response = tr.sent[0]
    assert "error" in response
    assert response["error"]["code"] == -32000
    assert response["error"]["message"] == "Internal error"
    assert "Network error" in response["error"]["data"]
    assert response["id"] == 1


@pytest.mark.asyncio
async def test_session_backend_docstring_examples():
    """Test the docstring examples in SessionBackend."""
    # First-Party
    from mcpgateway.cache.session_registry import SessionBackend

    # Test memory backend example
    backend = SessionBackend(backend="memory")
    assert backend._backend == "memory"
    assert backend._session_ttl == 3600

    # Test redis backend without URL
    try:
        backend = SessionBackend(backend="redis")
        assert False, "Should have raised ValueError"
    except ValueError as e:
        # In environments without redis installed, error message differs
        assert ("Redis backend requires redis_url" in str(e)) or ("redis package not installed" in str(e))

    # Test invalid backend
    try:
        backend = SessionBackend(backend="invalid")
        assert False, "Should have raised ValueError"
    except ValueError as e:
        assert "Invalid backend" in str(e)


@pytest.mark.asyncio
async def test_redis_initialize_error(monkeypatch):
    """Test Redis initialization with pubsub error."""
    mock_redis = MockRedis()

    class MockPubSubWithError:
        async def subscribe(self, channel):
            raise Exception("pubsub subscribe error")

    mock_redis.pubsub = lambda: MockPubSubWithError()

    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(backend="redis", redis_url="redis://localhost:6379")

        # This should raise the exception from pubsub.subscribe
        with pytest.raises(Exception, match="pubsub subscribe error"):
            await registry.initialize()


@pytest.mark.asyncio
async def test_memory_cleanup_task_error(monkeypatch, caplog):
    """Test memory cleanup task with transport error."""
    registry = SessionRegistry(backend="memory")

    # Mock transport that raises error on is_connected
    tr = FakeSSETransport("error_session")
    tr.is_connected = AsyncMock(side_effect=Exception("connection check error"))

    await registry.add_session("error_session", tr)

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    # Start cleanup task
    cleanup_task = asyncio.create_task(registry._memory_cleanup_task())

    # Let it run briefly
    await asyncio.sleep(0.01)

    # Cancel the task
    cleanup_task.cancel()
    try:
        await cleanup_task
    except asyncio.CancelledError:
        pass

    assert "Error checking session error_session" in caplog.text
    # Session should be removed due to error
    assert registry.get_session_sync("error_session") is None


@pytest.mark.asyncio
async def test_memory_cleanup_task_cancelled(monkeypatch, caplog):
    """Test memory cleanup task cancellation."""
    registry = SessionRegistry(backend="memory")

    caplog.set_level(logging.INFO, logger="mcpgateway.cache.session_registry")

    # Start cleanup task
    cleanup_task = asyncio.create_task(registry._memory_cleanup_task())

    # Let it start
    await asyncio.sleep(0.001)

    # Cancel the task
    cleanup_task.cancel()
    try:
        await cleanup_task
    except asyncio.CancelledError:
        pass

    assert "Memory cleanup task cancelled" in caplog.text


@pytest.mark.asyncio
async def test_refresh_redis_sessions(monkeypatch):
    """Test _refresh_redis_sessions method."""
    mock_redis = MockRedis()

    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(backend="redis", redis_url="redis://localhost:6379")
        await registry.initialize()

        # Add connected session
        tr1 = FakeSSETransport("connected_session", connected=True)
        await registry.add_session("connected_session", tr1)

        # Add disconnected session
        tr2 = FakeSSETransport("disconnected_session", connected=False)
        await registry.add_session("disconnected_session", tr2)

        await registry._refresh_redis_sessions()

        # Connected session should still exist
        assert registry.get_session_sync("connected_session") is tr1

        # Disconnected session should be removed
        assert registry.get_session_sync("disconnected_session") is None

        await registry.shutdown()


@pytest.mark.asyncio
async def test_refresh_redis_sessions_error(monkeypatch, caplog):
    """Test _refresh_redis_sessions with transport error."""
    mock_redis = MockRedis()

    monkeypatch.setattr("mcpgateway.cache.session_registry.REDIS_AVAILABLE", True)

    # Mock the shared Redis client factory to return our mock
    async def mock_get_redis_client():
        return mock_redis

    caplog.set_level(logging.ERROR, logger="mcpgateway.cache.session_registry")

    with patch("mcpgateway.cache.session_registry.get_redis_client", mock_get_redis_client):
        registry = SessionRegistry(backend="redis", redis_url="redis://localhost:6379")
        await registry.initialize()

        # Add session with transport that raises error
        tr = FakeSSETransport("error_session")
        tr.is_connected = AsyncMock(side_effect=Exception("connection error"))
        await registry.add_session("error_session", tr)

        await registry._refresh_redis_sessions()

        assert "Error refreshing session error_session" in caplog.text

        await registry.shutdown()


# --------------------------------------------------------------------------- #
# Respond Task Tracking and Cancellation Tests                                #
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_register_respond_task(registry: SessionRegistry):
    """Test registering a respond task for a session."""

    # Create a simple task
    async def dummy_task():
        await asyncio.sleep(10)

    task = asyncio.create_task(dummy_task())

    # Register it
    registry.register_respond_task("test_session", task)

    # Verify it's tracked
    assert "test_session" in registry._respond_tasks
    assert registry._respond_tasks["test_session"] is task

    # Cleanup
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass


@pytest.mark.asyncio
async def test_cancel_respond_task_success(registry: SessionRegistry):
    """Test successful cancellation of a respond task."""
    cancelled = asyncio.Event()

    async def cancellable_task():
        try:
            await asyncio.sleep(10)
        except asyncio.CancelledError:
            cancelled.set()
            raise

    task = asyncio.create_task(cancellable_task())
    await asyncio.sleep(0)  # Let task start executing
    registry.register_respond_task("cancel_test", task)

    # Cancel the task
    await registry._cancel_respond_task("cancel_test")

    # Verify task was cancelled and removed from tracking
    assert cancelled.is_set()
    assert "cancel_test" not in registry._respond_tasks


@pytest.mark.asyncio
async def test_cancel_respond_task_already_done(registry: SessionRegistry):
    """Test cancelling an already-completed task."""
    completed = asyncio.Event()

    async def quick_task():
        completed.set()
        return "done"

    task = asyncio.create_task(quick_task())
    await asyncio.sleep(0.01)  # Let task complete

    assert task.done()
    registry.register_respond_task("done_test", task)

    # Cancel should handle already-done task gracefully
    await registry._cancel_respond_task("done_test")

    # Verify removed from tracking
    assert "done_test" not in registry._respond_tasks


@pytest.mark.asyncio
async def test_cancel_respond_task_nonexistent(registry: SessionRegistry):
    """Test cancelling a task that doesn't exist."""
    # Should not raise any errors
    await registry._cancel_respond_task("nonexistent_session")


@pytest.mark.asyncio
async def test_cancel_respond_task_timeout_with_escalation(registry: SessionRegistry, caplog):
    """Test task cancellation timeout triggers escalation."""
    caplog.set_level(logging.WARNING, logger="mcpgateway.cache.session_registry")

    stop_flag = asyncio.Event()

    async def stubborn_task():
        # Use shield to protect against cancellation
        try:
            await asyncio.shield(asyncio.sleep(100))
        except asyncio.CancelledError:
            # Still ignore - wait until told to stop
            while not stop_flag.is_set():
                await asyncio.sleep(0.001)

    task = asyncio.create_task(stubborn_task())
    await asyncio.sleep(0)  # Let task start executing
    registry.register_respond_task("stubborn_test", task)

    # Cancel with very short timeout - will trigger escalation
    await registry._cancel_respond_task("stubborn_test", timeout=0.05)

    # After escalation, task should be removed from _respond_tasks
    # It may be in _stuck_tasks if retry also failed, or removed if retry succeeded
    assert "stubborn_test" not in registry._respond_tasks
    # Should see escalation message
    assert "escalating" in caplog.text

    # Cleanup - signal task to stop and clear any stuck_tasks
    stop_flag.set()
    task.cancel()
    try:
        await asyncio.wait_for(task, timeout=0.1)
    except (asyncio.CancelledError, asyncio.TimeoutError):
        pass
    registry._stuck_tasks.pop("stubborn_test", None)


@pytest.mark.asyncio
async def test_cancel_respond_task_escalation_calls_transport_disconnect(registry: SessionRegistry, caplog):
    """Test that escalation calls transport.disconnect() to unblock the task."""
    caplog.set_level(logging.DEBUG, logger="mcpgateway.cache.session_registry")

    stop_flag = asyncio.Event()

    async def stubborn_task():
        # Task that ignores cancellation until told to stop
        try:
            await asyncio.shield(asyncio.sleep(100))
        except asyncio.CancelledError:
            while not stop_flag.is_set():
                await asyncio.sleep(0.001)

    # Create a fake transport and add it to sessions
    transport = FakeSSETransport("escalation_test")
    await registry.add_session("escalation_test", transport)

    # Register the stubborn task
    task = asyncio.create_task(stubborn_task())
    await asyncio.sleep(0)  # Let task start executing
    registry.register_respond_task("escalation_test", task)

    # Cancel with very short timeout - will trigger escalation
    await registry._cancel_respond_task("escalation_test", timeout=0.05)

    # Verify transport.disconnect() was called during escalation
    assert transport.disconnect_called, "transport.disconnect() should be called during escalation"

    # Verify escalation message in logs
    assert "Force-disconnected transport" in caplog.text or "escalating" in caplog.text

    # Cleanup
    stop_flag.set()
    task.cancel()
    try:
        await asyncio.wait_for(task, timeout=0.1)
    except (asyncio.CancelledError, asyncio.TimeoutError):
        pass
    registry._stuck_tasks.pop("escalation_test", None)


@pytest.mark.asyncio
async def test_cancel_respond_task_with_exception(registry: SessionRegistry, caplog):
    """Test cancelling a task that failed with an exception."""
    caplog.set_level(logging.WARNING, logger="mcpgateway.cache.session_registry")

    async def failing_task():
        raise ValueError("Task failed!")

    task = asyncio.create_task(failing_task())
    await asyncio.sleep(0.01)  # Let task fail

    assert task.done()
    registry.register_respond_task("fail_test", task)

    # Cancel should handle failed task and log warning
    await registry._cancel_respond_task("fail_test")

    # Verify removed from tracking
    assert "fail_test" not in registry._respond_tasks
    assert "failed with" in caplog.text


@pytest.mark.asyncio
async def test_remove_session_cancels_respond_task(registry: SessionRegistry):
    """Test that remove_session cancels the respond task first."""
    cancelled = asyncio.Event()

    async def tracked_task():
        try:
            await asyncio.sleep(10)
        except asyncio.CancelledError:
            cancelled.set()
            raise

    # Add session with transport
    tr = FakeSSETransport("session_with_task")
    await registry.add_session("session_with_task", tr)

    # Register respond task
    task = asyncio.create_task(tracked_task())
    await asyncio.sleep(0)  # Let task start executing
    registry.register_respond_task("session_with_task", task)

    # Remove session should cancel task
    await registry.remove_session("session_with_task")

    # Verify task was cancelled
    assert cancelled.is_set()
    assert "session_with_task" not in registry._respond_tasks


@pytest.mark.asyncio
async def test_shutdown_cancels_all_respond_tasks(registry: SessionRegistry):
    """Test that shutdown cancels all respond tasks."""
    cancelled_count = 0

    async def make_task(name):
        nonlocal cancelled_count
        try:
            await asyncio.sleep(10)
        except asyncio.CancelledError:
            cancelled_count += 1
            raise

    # Register multiple tasks
    tasks = []
    for i in range(3):
        task = asyncio.create_task(make_task(f"task_{i}"))
        registry.register_respond_task(f"session_{i}", task)
        tasks.append(task)

    await asyncio.sleep(0)  # Let all tasks start executing

    # Shutdown should cancel all
    await registry.shutdown()

    # Verify all cancelled
    assert cancelled_count == 3
    assert len(registry._respond_tasks) == 0


@pytest.mark.asyncio
async def test_register_respond_task_overwrites_previous():
    """Test that registering a new task overwrites the previous one."""
    registry = SessionRegistry(backend="memory")
    await registry.initialize()

    async def task1():
        await asyncio.sleep(10)

    async def task2():
        await asyncio.sleep(10)

    t1 = asyncio.create_task(task1())
    t2 = asyncio.create_task(task2())

    registry.register_respond_task("overwrite_test", t1)
    registry.register_respond_task("overwrite_test", t2)

    # Only t2 should be tracked
    assert registry._respond_tasks["overwrite_test"] is t2

    # Cleanup
    t1.cancel()
    t2.cancel()
    try:
        await t1
    except asyncio.CancelledError:
        pass
    try:
        await t2
    except asyncio.CancelledError:
        pass

    await registry.shutdown()


@pytest.mark.asyncio
async def test_stuck_task_reaper_cleans_completed_tasks(caplog):
    """Test that the stuck task reaper cleans up completed tasks."""
    caplog.set_level(logging.INFO, logger="mcpgateway.cache.session_registry")

    registry = SessionRegistry(backend="memory")
    await registry.initialize()

    # Manually add a completed task to _stuck_tasks
    completed_event = asyncio.Event()

    async def quick_task():
        completed_event.set()
        return "done"

    task = asyncio.create_task(quick_task())
    await asyncio.sleep(0.01)  # Let task complete
    assert task.done()

    registry._stuck_tasks["completed_stuck"] = task

    # Manually call the reaper logic (don't wait for the loop)
    # Simulate one iteration of the reaper
    for session_id, t in list(registry._stuck_tasks.items()):
        if t.done():
            registry._stuck_tasks.pop(session_id, None)

    # Verify task was removed
    assert "completed_stuck" not in registry._stuck_tasks

    await registry.shutdown()


@pytest.mark.asyncio
async def test_shutdown_cancels_stuck_task_reaper():
    """Test that shutdown properly cancels the stuck task reaper."""
    registry = SessionRegistry(backend="memory")
    await registry.initialize()

    # Verify reaper was started
    assert registry._stuck_task_reaper is not None
    assert not registry._stuck_task_reaper.done()

    await registry.shutdown()

    # Verify reaper was cancelled
    assert registry._stuck_task_reaper.done()


if __name__ == "__main__":
    # Allow running tests directly
    pytest.main([__file__, "-v"])
