from typing import Self, AsyncGenerator

from ariadne.types import Resolver

from api.broadcast import broadcast
from api_graphql.abc.field_builder_abc import FieldBuilderABC
from api_graphql.field.subscription_field import SubscriptionField


class SubscriptionFieldBuilder(FieldBuilderABC):

    def __init__(self, name: str):
        FieldBuilderABC.__init__(self, name)

        self._resolver = None
        self._generator = None

    def with_resolver(self, resolver: Resolver) -> Self:
        assert resolver is not None, "resolver cannot be None"
        self._resolver = resolver
        return self

    def with_generator(self, generator: Resolver) -> Self:
        assert generator is not None, "generator cannot be None"
        self._generator = generator
        return self

    def build(self) -> SubscriptionField:
        assert self._resolver is not None, "resolver cannot be None"
        if self._generator is None:

            async def generator(*args, **kwargs) -> AsyncGenerator[str, None]:
                async with broadcast.subscribe(channel=self._name) as subscriber:
                    async for message in subscriber:
                        yield message

            self._generator = generator

        return SubscriptionField(
            self._name,
            self._require_any_permission,
            self._require_any,
            self._public,
            self._resolver,
            self._generator,
        )
