# retoor <retoor@molodetz.nl>

from __future__ import annotations

import logging

logger = logging.getLogger(__name__)


def topic_matches(pattern: str, topic: str) -> bool:
    if pattern == topic or pattern == "*":
        return True
    if pattern.endswith(".*"):
        prefix = pattern[:-1]
        return topic == pattern[:-2] or topic.startswith(prefix)
    return False


def _is_pattern(pattern: str) -> bool:
    return pattern == "*" or pattern.endswith(".*")


class PubSubHub:
    def __init__(self) -> None:
        self._subscriptions: dict[str, set] = {}
        self._patterns: dict[str, set] = {}

    def _store_for(self, pattern: str) -> dict:
        return self._patterns if _is_pattern(pattern) else self._subscriptions

    def subscribe(self, pattern: str, socket) -> None:
        self._store_for(pattern).setdefault(pattern, set()).add(socket)

    def unsubscribe(self, pattern: str, socket) -> None:
        store = self._store_for(pattern)
        sockets = store.get(pattern)
        if not sockets:
            return
        sockets.discard(socket)
        if not sockets:
            store.pop(pattern, None)

    def drop_socket(self, socket) -> None:
        for store in (self._subscriptions, self._patterns):
            for pattern in list(store):
                sockets = store.get(pattern)
                if not sockets:
                    continue
                sockets.discard(socket)
                if not sockets:
                    store.pop(pattern, None)

    def _targets(self, topic: str) -> set:
        targets: set = set(self._subscriptions.get(topic, ()))
        for pattern, sockets in self._patterns.items():
            if topic_matches(pattern, topic):
                targets.update(sockets)
        return targets

    async def publish(self, topic: str, frame: dict) -> int:
        targets = self._targets(topic)
        delivered = 0
        for socket in targets:
            try:
                await socket.send_json(frame)
                delivered += 1
            except Exception:
                logger.debug("pubsub dropping dead socket for %s", topic)
        return delivered

    def topics(self) -> list[dict]:
        merged = {**self._subscriptions, **self._patterns}
        return [
            {"topic": pattern, "subscribers": len(sockets)}
            for pattern, sockets in sorted(merged.items())
        ]


pubsub = PubSubHub()
