"""Typed in-memory subscriber registry by subscription category.

Queue execution subscriptions use job_id under QUEUE_JOB; WebSocket delivery uses ws_id under
WS_DELIVERY; transfer session events use transfer_id under TRANSFER_SESSION.
Authorization lives outside this module (e.g. ws_identity).
Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from enum import Enum
from typing import Dict, List, Set


class SubscriptionCategory(Enum):
    """Canonical subscription domain discriminator."""

    QUEUE_JOB = "queue_job"
    WS_DELIVERY = "ws_delivery"
    TRANSFER_SESSION = "transfer_session"


class Subscriber(ABC):
    """Abstract interface for sending a message to one connection (e.g. WebSocket)."""

    @abstractmethod
    async def send(self, payload: dict) -> None:
        """Send payload to this subscriber. May raise on closed connection."""
        raise NotImplementedError(
            "Subscriber.send must be implemented by concrete class"
        )


def _normalize_key(key: str) -> str:
    """Return stripped key; caller checks for empty."""
    return str(key).strip()


class SubscriberRegistry:
    """
    Typed dual-index registry: (key, category) -> subscribers, subscriber -> (key, category).

    Thread/async safety: all operations use a single asyncio.Lock.
    Queue polling must use get_subscribed_keys_for_queue_poll() only.
    """

    def __init__(self) -> None:
        """Initialize empty registry with a lock."""
        self._lock = asyncio.Lock()
        self._subscriptions: Dict[tuple[str, SubscriptionCategory], Set[Subscriber]] = (
            {}
        )
        self._subscriber_index: Dict[
            Subscriber, Set[tuple[str, SubscriptionCategory]]
        ] = {}

    async def subscribe(
        self,
        key: str,
        category: SubscriptionCategory,
        subscriber: Subscriber,
    ) -> None:
        """Register subscription for (key, category)."""
        k = _normalize_key(key)
        if k == "":
            return
        async with self._lock:
            self._subscriptions.setdefault((k, category), set()).add(subscriber)
            self._subscriber_index.setdefault(subscriber, set()).add((k, category))

    async def unsubscribe(
        self,
        key: str,
        category: SubscriptionCategory,
        subscriber: Subscriber,
    ) -> None:
        """Remove one binding for (key, category)."""
        k = _normalize_key(key)
        if k == "":
            return
        async with self._lock:
            subs = self._subscriptions.get((k, category))
            if subs:
                subs.discard(subscriber)
                if not subs:
                    del self._subscriptions[(k, category)]
            pairs = self._subscriber_index.get(subscriber)
            if pairs:
                pairs.discard((k, category))
                if not pairs:
                    del self._subscriber_index[subscriber]

    async def unsubscribe_all_for(self, subscriber: Subscriber) -> None:
        """Remove this subscriber from all (key, category) pairs."""
        async with self._lock:
            pairs = self._subscriber_index.pop(subscriber, None)
            if pairs is None:
                return
            for k, cat in pairs:
                subs = self._subscriptions.get((k, cat))
                if subs:
                    subs.discard(subscriber)
                    if not subs:
                        del self._subscriptions[(k, cat)]

    async def get_subscribers(
        self, key: str, category: SubscriptionCategory
    ) -> List[Subscriber]:
        """Return a snapshot of subscribers for this key+category."""
        k = _normalize_key(key)
        if k == "":
            return []
        async with self._lock:
            subs = self._subscriptions.get((k, category))
            return list(subs) if subs else []

    async def get_subscribers_for_delivery(self, ws_id: str) -> List[Subscriber]:
        """Return subscribers for ws_id under WS_DELIVERY only."""
        return await self.get_subscribers(ws_id, SubscriptionCategory.WS_DELIVERY)

    async def get_subscribed_keys_for_queue_poll(self) -> Set[str]:
        """Return keys that have at least one QUEUE_JOB subscriber."""
        async with self._lock:
            return {
                k
                for (k, cat) in self._subscriptions
                if cat == SubscriptionCategory.QUEUE_JOB
                and self._subscriptions[(k, cat)]
            }

    async def has_subscribers(self, key: str, category: SubscriptionCategory) -> bool:
        """Return True if there is at least one subscriber for (key, category)."""
        k = _normalize_key(key)
        if k == "":
            return False
        async with self._lock:
            subs = self._subscriptions.get((k, category))
            return bool(subs)
