"""Tests for typed SubscriberRegistry (queue_job vs ws_delivery).

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

from __future__ import annotations

import asyncio

import pytest

from mcp_proxy_adapter.core.job_push.subscriptions import (
    Subscriber,
    SubscriberRegistry,
    SubscriptionCategory,
)


class DummySubscriber(Subscriber):
    """No-op subscriber for registry tests."""

    async def send(self, payload: dict) -> None:
        """No-op send."""
        return None


@pytest.mark.asyncio
async def test_queue_poll_keys_exclude_ws_only() -> None:
    """get_subscribed_keys_for_queue_poll excludes WS_DELIVERY-only keys."""
    registry = SubscriberRegistry()
    sub = DummySubscriber()
    await registry.subscribe("ws-9", SubscriptionCategory.WS_DELIVERY, sub)
    keys = await registry.get_subscribed_keys_for_queue_poll()
    assert "ws-9" not in keys
    assert keys == set()


@pytest.mark.asyncio
async def test_queue_poll_keys_include_queue_job() -> None:
    """get_subscribed_keys_for_queue_poll includes QUEUE_JOB keys."""
    registry = SubscriberRegistry()
    sub = DummySubscriber()
    await registry.subscribe("job-1", SubscriptionCategory.QUEUE_JOB, sub)
    keys = await registry.get_subscribed_keys_for_queue_poll()
    assert keys == {"job-1"}


@pytest.mark.asyncio
async def test_delivery_subscribers_isolated() -> None:
    """QUEUE_JOB and WS_DELIVERY subscribers for same key are isolated."""
    registry = SubscriberRegistry()
    sub_a = DummySubscriber()
    sub_b = DummySubscriber()
    await registry.subscribe("same-key", SubscriptionCategory.QUEUE_JOB, sub_a)
    await registry.subscribe("same-key", SubscriptionCategory.WS_DELIVERY, sub_b)
    q = await registry.get_subscribers("same-key", SubscriptionCategory.QUEUE_JOB)
    w = await registry.get_subscribers("same-key", SubscriptionCategory.WS_DELIVERY)
    assert q == [sub_a]
    assert w == [sub_b]


@pytest.mark.asyncio
async def test_unsubscribe_all_removes_both_domains() -> None:
    """unsubscribe_all_for removes subscriber from all (key, category) pairs."""
    registry = SubscriberRegistry()
    sub = DummySubscriber()
    await registry.subscribe("j1", SubscriptionCategory.QUEUE_JOB, sub)
    await registry.subscribe("w1", SubscriptionCategory.WS_DELIVERY, sub)
    await registry.unsubscribe_all_for(sub)
    assert await registry.get_subscribers("j1", SubscriptionCategory.QUEUE_JOB) == []
    assert await registry.get_subscribers("w1", SubscriptionCategory.WS_DELIVERY) == []
    assert await registry.get_subscribed_keys_for_queue_poll() == set()


@pytest.mark.asyncio
async def test_concurrent_subscribe_unsubscribe() -> None:
    """Concurrent subscribe/unsubscribe does not corrupt registry."""
    registry = SubscriberRegistry()

    async def worker(i: int) -> None:
        sub = DummySubscriber()
        await registry.subscribe(f"k{i}", SubscriptionCategory.QUEUE_JOB, sub)
        await registry.unsubscribe(f"k{i}", SubscriptionCategory.QUEUE_JOB, sub)

    await asyncio.gather(*[worker(n) for n in range(20)])
    assert await registry.get_subscribed_keys_for_queue_poll() == set()
