"""
In-memory registry: job_id -> set of subscribers.

Used for pushing job events to WebSocket (or other) connections.
No auth in this module; authorization is done in api/ws_identity before subscribe.
See docs/plans/bidirectional_job_push/03_subscriber_registry.md.

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

from __future__ import annotations

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


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")


class SubscriberRegistry:
    """
    Registry mapping job_id -> set of subscribers.

    Thread/async safety: all operations use a single asyncio.Lock.
    Lifecycle: instance is created in lifespan and stored in app.state.job_push_registry.
    """

    def __init__(self) -> None:
        """Initialize empty registry with a lock."""
        self._lock = asyncio.Lock()
        self._job_to_subscribers: Dict[str, Set[object]] = {}
        self._subscriber_to_jobs: Dict[object, Set[str]] = {}

    async def subscribe(self, job_id: str, subscriber: Subscriber) -> None:
        """Add subscriber for the given job_id."""
        async with self._lock:
            self._job_to_subscribers.setdefault(job_id, set()).add(subscriber)
            self._subscriber_to_jobs.setdefault(subscriber, set()).add(job_id)

    async def unsubscribe(self, job_id: str, subscriber: Subscriber) -> None:
        """Remove subscriber from the given job_id."""
        async with self._lock:
            subs = self._job_to_subscribers.get(job_id)
            if subs:
                subs.discard(subscriber)
                if not subs:
                    del self._job_to_subscribers[job_id]
            jobs = self._subscriber_to_jobs.get(subscriber)
            if jobs:
                jobs.discard(job_id)
                if not jobs:
                    del self._subscriber_to_jobs[subscriber]

    async def unsubscribe_all_for(self, subscriber: Subscriber) -> None:
        """Remove this subscriber from all job_ids (e.g. on disconnect)."""
        async with self._lock:
            jobs = self._subscriber_to_jobs.pop(subscriber, None)
            if jobs:
                for jid in jobs:
                    subs = self._job_to_subscribers.get(jid)
                    if subs:
                        subs.discard(subscriber)
                        if not subs:
                            del self._job_to_subscribers[jid]

    async def get_subscribers(self, job_id: str) -> List[Subscriber]:
        """Return a snapshot of subscribers for this job_id."""
        async with self._lock:
            subs = self._job_to_subscribers.get(job_id)
            return list(subs) if subs else []

    async def get_subscribed_job_ids(self) -> Set[str]:
        """Return the set of job_ids that have at least one subscriber (for poll loop)."""
        async with self._lock:
            return set(self._job_to_subscribers.keys())
