"""Tests for typed WebSocket subscribe/unsubscribe protocol.

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,
    SUBSCRIPTION_TYPE_QUEUE_JOB,
    SUBSCRIPTION_TYPE_WS_DELIVERY,
)
from mcp_proxy_adapter.core.job_push.subscriptions import SubscriptionCategory


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


@pytest.mark.asyncio
async def test_queue_subscribe_uses_category() -> None:
    """Subscribe with subscription_type queue_job calls registry with QUEUE_JOB."""
    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_QUEUE_JOB,
                    "job_id": "job-1",
                }
            ),
            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)
    assert registry.subscribe.await_count >= 1
    call_args = registry.subscribe.await_args
    assert call_args[0][0] == "job-1"
    assert call_args[0][1] == SubscriptionCategory.QUEUE_JOB


@pytest.mark.asyncio
async def test_ws_subscribe_triggers_replay_flush() -> None:
    """Subscribe with ws_delivery triggers flush_replay_to_subscriber."""
    app = _make_app_with_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_WS_DELIVERY,
                    "ws_id": "ws-456",
                }
            ),
            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)
    manager = app.state.ws_delivery_manager
    assert manager.flush_replay_to_subscriber.await_count >= 1
    call_args = manager.flush_replay_to_subscriber.await_args[0]
    assert call_args[0] == "ws-456"


@pytest.mark.asyncio
async def test_unknown_type_errors() -> None:
    """Unknown subscription_type yields subscribe_error."""
    app = _make_app_with_registry()
    ws = MagicMock()
    ws.accept = AsyncMock()
    ws.scope = {"type": "websocket", "headers": []}
    ws.receive_text = AsyncMock(
        side_effect=[
            json.dumps(
                {
                    "action": ACTION_SUBSCRIBE,
                    "subscription_type": "nope",
                }
            ),
            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)
    assert ws.send_json.called
    payload = ws.send_json.call_args[0][0]
    assert payload.get("action") == ACTION_SUBSCRIBE_ERROR
    assert "unknown" in str(payload.get("error", "")).lower()
