"""Tests for WebSocket subscribe/unsubscribe with subscription_type transfer_session.

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import FastAPI
from fastapi import WebSocketDisconnect

from mcp_proxy_adapter.api.ws_endpoint import handle_websocket_job_push
from mcp_proxy_adapter.core.job_push.events import (
    ACTION_SUBSCRIBE,
    ACTION_SUBSCRIBE_ERROR,
    ACTION_UNSUBSCRIBE,
    SUBSCRIPTION_TYPE_TRANSFER_SESSION,
)
from mcp_proxy_adapter.core.job_push.subscriptions import (
    Subscriber,
    SubscriberRegistry,
    SubscriptionCategory,
)
from mcp_proxy_adapter.transfer.errors import (
    TransferExpiredError,
    TransferSessionNotFoundError,
)


def _make_app_with_registry() -> FastAPI:
    """Build FastAPI app with mocked registry and ws_delivery_manager."""
    app = FastAPI()
    registry = MagicMock()
    registry.subscribe = AsyncMock()
    registry.unsubscribe = AsyncMock()
    registry.unsubscribe_all_for = AsyncMock()
    app.state.job_push_registry = registry
    manager = MagicMock()
    manager.flush_replay_to_subscriber = AsyncMock()
    app.state.ws_delivery_manager = manager
    return app


class _DummySubscriber(Subscriber):
    async def send(self, payload: dict) -> None:
        return None


@pytest.mark.asyncio
async def test_transfer_session_subscribe_uses_category() -> None:
    app = _make_app_with_registry()
    registry = app.state.job_push_registry
    ws = MagicMock()
    ws.accept = AsyncMock()
    ws.scope = {"type": "websocket", "headers": []}
    ws.receive_text = AsyncMock(
        side_effect=[
            json.dumps(
                {
                    "action": ACTION_SUBSCRIBE,
                    "subscription_type": SUBSCRIPTION_TYPE_TRANSFER_SESSION,
                    "transfer_id": "t-1",
                }
            ),
            WebSocketDisconnect(),
        ]
    )
    ws.send_json = AsyncMock()
    with patch(
        "mcp_proxy_adapter.api.ws_endpoint.is_allowed_to_subscribe",
        return_value=True,
    ), patch("mcp_proxy_adapter.api.handlers.get_transfer_store") as get_store:
        store = MagicMock()
        store.get_session.return_value = MagicMock()
        get_store.return_value = store
        await handle_websocket_job_push(ws, app)
    assert registry.subscribe.await_count >= 1
    call_args = registry.subscribe.await_args
    assert call_args[0][0] == "t-1"
    assert call_args[0][1] == SubscriptionCategory.TRANSFER_SESSION


@pytest.mark.asyncio
async def test_transfer_session_subscribe_not_found_sends_subscribe_error() -> None:
    app = _make_app_with_registry()
    registry = app.state.job_push_registry
    ws = MagicMock()
    ws.accept = AsyncMock()
    ws.scope = {"type": "websocket", "headers": []}
    ws.receive_text = AsyncMock(
        side_effect=[
            json.dumps(
                {
                    "action": ACTION_SUBSCRIBE,
                    "subscription_type": SUBSCRIPTION_TYPE_TRANSFER_SESSION,
                    "transfer_id": "missing",
                }
            ),
            WebSocketDisconnect(),
        ]
    )
    ws.send_json = AsyncMock()
    with patch(
        "mcp_proxy_adapter.api.ws_endpoint.is_allowed_to_subscribe",
        return_value=True,
    ), patch("mcp_proxy_adapter.api.handlers.get_transfer_store") as get_store:
        store = MagicMock()
        store.get_session.side_effect = TransferSessionNotFoundError(
            "Transfer session not found", transfer_id="missing"
        )
        get_store.return_value = store
        await handle_websocket_job_push(ws, app)
    assert ws.send_json.called
    payload = ws.send_json.call_args[0][0]
    assert payload["action"] == ACTION_SUBSCRIBE_ERROR
    assert payload["error"] == "transfer_id not found"
    assert payload["transfer_id"] == "missing"
    assert registry.subscribe.await_count == 0


@pytest.mark.asyncio
async def test_transfer_session_subscribe_expired_sends_subscribe_error() -> None:
    app = _make_app_with_registry()
    registry = app.state.job_push_registry
    ws = MagicMock()
    ws.accept = AsyncMock()
    ws.scope = {"type": "websocket", "headers": []}
    ws.receive_text = AsyncMock(
        side_effect=[
            json.dumps(
                {
                    "action": ACTION_SUBSCRIBE,
                    "subscription_type": SUBSCRIPTION_TYPE_TRANSFER_SESSION,
                    "transfer_id": "t-xp",
                }
            ),
            WebSocketDisconnect(),
        ]
    )
    ws.send_json = AsyncMock()
    with patch(
        "mcp_proxy_adapter.api.ws_endpoint.is_allowed_to_subscribe",
        return_value=True,
    ), patch("mcp_proxy_adapter.api.handlers.get_transfer_store") as get_store:
        store = MagicMock()
        store.get_session.side_effect = TransferExpiredError(
            "Transfer session expired", transfer_id="t-xp"
        )
        get_store.return_value = store
        await handle_websocket_job_push(ws, app)
    payload = ws.send_json.call_args[0][0]
    assert payload["action"] == ACTION_SUBSCRIBE_ERROR
    assert payload["error"] == "transfer session expired"
    assert payload["transfer_id"] == "t-xp"
    assert registry.subscribe.await_count == 0


@pytest.mark.asyncio
async def test_transfer_session_subscribe_missing_transfer_id_errors() -> None:
    app = _make_app_with_registry()
    registry = app.state.job_push_registry
    ws = MagicMock()
    ws.accept = AsyncMock()
    ws.scope = {"type": "websocket", "headers": []}
    ws.receive_text = AsyncMock(
        side_effect=[
            json.dumps(
                {
                    "action": ACTION_SUBSCRIBE,
                    "subscription_type": SUBSCRIPTION_TYPE_TRANSFER_SESSION,
                }
            ),
            WebSocketDisconnect(),
        ]
    )
    ws.send_json = AsyncMock()
    with patch(
        "mcp_proxy_adapter.api.ws_endpoint.is_allowed_to_subscribe",
        return_value=True,
    ):
        await handle_websocket_job_push(ws, app)
    payload = ws.send_json.call_args[0][0]
    assert payload["action"] == ACTION_SUBSCRIBE_ERROR
    assert payload["error"] == "transfer_id required for transfer_session"
    assert registry.subscribe.await_count == 0


@pytest.mark.asyncio
async def test_transfer_session_unsubscribe_removes_subscription() -> None:
    app = FastAPI()
    app.state.job_push_registry = SubscriberRegistry()
    app.state.ws_delivery_manager = MagicMock()
    app.state.ws_delivery_manager.flush_replay_to_subscriber = AsyncMock()
    registry = app.state.job_push_registry
    ws = MagicMock()
    ws.accept = AsyncMock()
    ws.scope = {"type": "websocket", "headers": []}
    ws.receive_text = AsyncMock(
        side_effect=[
            json.dumps(
                {
                    "action": ACTION_SUBSCRIBE,
                    "subscription_type": SUBSCRIPTION_TYPE_TRANSFER_SESSION,
                    "transfer_id": "t-unsub",
                }
            ),
            json.dumps(
                {
                    "action": ACTION_UNSUBSCRIBE,
                    "subscription_type": SUBSCRIPTION_TYPE_TRANSFER_SESSION,
                    "transfer_id": "t-unsub",
                }
            ),
            WebSocketDisconnect(),
        ]
    )
    ws.send_json = AsyncMock()
    with patch(
        "mcp_proxy_adapter.api.ws_endpoint.is_allowed_to_subscribe",
        return_value=True,
    ), patch("mcp_proxy_adapter.api.handlers.get_transfer_store") as get_store:
        store = MagicMock()
        store.get_session.return_value = MagicMock()
        get_store.return_value = store
        await handle_websocket_job_push(ws, app)
    subs = await registry.get_subscribers(
        "t-unsub", SubscriptionCategory.TRANSFER_SESSION
    )
    assert subs == []


@pytest.mark.asyncio
async def test_registry_transfer_session_isolated_from_queue_job_same_key() -> None:
    registry = SubscriberRegistry()
    sub_a = _DummySubscriber()
    sub_b = _DummySubscriber()
    await registry.subscribe("same-key", SubscriptionCategory.QUEUE_JOB, sub_a)
    await registry.subscribe("same-key", SubscriptionCategory.TRANSFER_SESSION, sub_b)
    q = await registry.get_subscribers("same-key", SubscriptionCategory.QUEUE_JOB)
    t = await registry.get_subscribers(
        "same-key", SubscriptionCategory.TRANSFER_SESSION
    )
    assert q == [sub_a]
    assert t == [sub_b]
    assert sub_b not in q
    assert sub_a not in t
