# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/015_FastKafka.ipynb.

# %% auto 0
__all__ = ['logger', 'FastKafka', 'AwaitedMock']

# %% ../../nbs/015_FastKafka.ipynb 1
import asyncio
import functools
import inspect
import json
import types
from asyncio import iscoroutinefunction  # do not use the version from inspect
from collections import namedtuple
from contextlib import asynccontextmanager
from datetime import datetime, timedelta
from functools import wraps
from inspect import signature
from pathlib import Path
from typing import *
from unittest.mock import AsyncMock, MagicMock

import anyio
from pydantic import BaseModel

from fastkafka._components.aiokafka_consumer_loop import (
    aiokafka_consumer_loop,
    sanitize_kafka_config,
)
from fastkafka._components.asyncapi import (
    ConsumeCallable,
    ContactInfo,
    KafkaBroker,
    KafkaBrokers,
    KafkaServiceInfo,
    export_async_spec,
)

import fastkafka._aiokafka_imports
from .._aiokafka_imports import AIOKafkaConsumer, AIOKafkaProducer
from .._components.benchmarking import _benchmark
from .._components.logger import get_logger
from .._components.meta import delegates, export, filter_using_signature, patch
from .._components.producer_decorator import ProduceCallable, producer_decorator
from .._components.task_streaming import StreamExecutor
from .._components.helpers import remove_suffix

# %% ../../nbs/015_FastKafka.ipynb 2
if TYPE_CHECKING:
    from fastapi import FastAPI

# %% ../../nbs/015_FastKafka.ipynb 4
logger = get_logger(__name__)

# %% ../../nbs/015_FastKafka.ipynb 9
@delegates(fastkafka._aiokafka_imports.AIOKafkaConsumer, but=["bootstrap_servers"])
@delegates(
    fastkafka._aiokafka_imports.AIOKafkaProducer, but=["bootstrap_servers"], keep=True
)
def _get_kafka_config(
    bootstrap_servers_id: str = "localhost",
    **kwargs: Any,
) -> Dict[str, Any]:
    """Get kafka config"""
    allowed_keys = set(signature(_get_kafka_config).parameters.keys())
    if not set(kwargs.keys()) <= allowed_keys:
        unallowed_keys = ", ".join(
            sorted([f"'{x}'" for x in set(kwargs.keys()).difference(allowed_keys)])
        )
        raise ValueError(f"Unallowed key arguments passed: {unallowed_keys}")
    retval = kwargs.copy()

    # todo: check this values
    config_defaults = {
        "bootstrap_servers_id": bootstrap_servers_id,
        "auto_offset_reset": "earliest",
        "max_poll_records": 100,
    }
    for key, value in config_defaults.items():
        if key not in retval:
            retval[key] = value

    return retval

# %% ../../nbs/015_FastKafka.ipynb 12
def _get_kafka_brokers(
    kafka_brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = None
) -> KafkaBrokers:
    """Get Kafka brokers

    Args:
        kafka_brokers: Kafka brokers

    """
    if kafka_brokers is None:
        retval: KafkaBrokers = KafkaBrokers(
            brokers={
                "localhost": KafkaBroker(  # type: ignore
                    url="https://localhost",
                    description="Local (dev) Kafka broker",
                    port="9092",
                    grouping="localhost",
                )
            }
        )
    else:
        if isinstance(kafka_brokers, KafkaBrokers):
            return kafka_brokers

        retval = KafkaBrokers(
            brokers={
                k: (
                    [
                        KafkaBroker.model_validate_json(
                            unwrapped_v.model_dump_json()
                            if hasattr(unwrapped_v, "model_dump_json")
                            else json.dumps(unwrapped_v)
                        )
                        for unwrapped_v in v
                    ]
                    if isinstance(v, list)
                    else KafkaBroker.model_validate_json(
                        v.model_dump_json()
                        if hasattr(v, "model_dump_json")
                        else json.dumps(v)
                    )
                )
                for k, v in kafka_brokers.items()
            }
        )

    return retval

# %% ../../nbs/015_FastKafka.ipynb 14
def _get_broker_addr_list(
    brokers: Union[List[KafkaBroker], KafkaBroker]
) -> Union[str, List[str]]:
    if isinstance(brokers, list):
        return [f"{broker.url}:{broker.port}" for broker in brokers]
    else:
        return f"{brokers.url}:{brokers.port}"

# %% ../../nbs/015_FastKafka.ipynb 16
def _get_topic_name(
    topic_callable: Union[ConsumeCallable, ProduceCallable], prefix: str = "on_"
) -> str:
    """Get topic name
    Args:
        topic_callable: a function
        prefix: prefix of the name of the function followed by the topic name

    Returns:
        The name of the topic
    """
    topic = topic_callable.__name__
    if not topic.startswith(prefix) or len(topic) <= len(prefix):
        raise ValueError(f"Function name '{topic}' must start with {prefix}")
    topic = topic[len(prefix) :]

    return topic

# %% ../../nbs/015_FastKafka.ipynb 18
def _get_contact_info(
    name: str = "Author",
    url: str = "https://www.google.com",
    email: str = "noreply@gmail.com",
) -> ContactInfo:
    return ContactInfo(name=name, url=url, email=email)  # type: ignore

# %% ../../nbs/015_FastKafka.ipynb 20
I = TypeVar("I", bound=BaseModel)
O = TypeVar("O", BaseModel, Awaitable[BaseModel])

F = TypeVar("F", bound=Callable)

# %% ../../nbs/015_FastKafka.ipynb 21
@export("fastkafka")
class FastKafka:
    @delegates(_get_kafka_config)
    def __init__(
        self,
        *,
        title: Optional[str] = None,
        description: Optional[str] = None,
        version: Optional[str] = None,
        contact: Optional[Dict[str, str]] = None,
        kafka_brokers: Optional[Dict[str, Any]] = None,
        root_path: Optional[Union[Path, str]] = None,
        lifespan: Optional[Callable[["FastKafka"], AsyncContextManager[None]]] = None,
        **kwargs: Any,
    ):
        """Creates FastKafka application

        Args:
            title: optional title for the documentation. If None,
                the title will be set to empty string
            description: optional description for the documentation. If
                None, the description will be set to empty string
            version: optional version for the documentation. If None,
                the version will be set to empty string
            contact: optional contact for the documentation. If None, the
                contact will be set to placeholder values:
                name='Author' url=HttpUrl('https://www.google.com', ) email='noreply@gmail.com'
            kafka_brokers: dictionary describing kafka brokers used for setting
                the bootstrap server when running the applicationa and for
                generating documentation. Defaults to
                    {
                        "localhost": {
                            "url": "localhost",
                            "description": "local kafka broker",
                            "port": "9092",
                        }
                    }
            root_path: path to where documentation will be created
            lifespan: asynccontextmanager that is used for setting lifespan hooks.
                __aenter__ is called before app start and __aexit__ after app stop.
                The lifespan is called whe application is started as async context
                manager, e.g.:`async with kafka_app...`

        """

        # this is needed for documentation generation
        self._title = title if title is not None else ""
        self._description = description if description is not None else ""
        self._version = version if version is not None else ""
        if contact is not None:
            self._contact_info = _get_contact_info(**contact)
        else:
            self._contact_info = _get_contact_info()

        self._kafka_service_info = KafkaServiceInfo(
            title=self._title,
            version=self._version,
            description=self._description,
            contact=self._contact_info,
        )

        if kafka_brokers is None:
            kafka_brokers = {
                "localhost": {
                    "url": "localhost",
                    "description": "local kafka broker",
                    "port": "9092",
                }
            }

        self._kafka_brokers = _get_kafka_brokers(kafka_brokers)

        self._override_brokers: List[KafkaBrokers] = []

        self._root_path = Path(".") if root_path is None else Path(root_path)
        self._root_path.mkdir(exist_ok=True, parents=True)

        self._asyncapi_path = self._root_path / "asyncapi"

        # this is used as default parameters for creating AIOProducer and AIOConsumer objects
        self._kafka_config = _get_kafka_config(**kwargs)

        #
        self._consumers_store: Dict[
            str,
            Tuple[
                ConsumeCallable,
                Callable[[bytes, Type[BaseModel]], Any],
                Union[str, StreamExecutor, None],
                Optional[KafkaBrokers],
                Dict[str, Any],
            ],
        ] = {}

        self._producers_store: Dict[  # type: ignore
            str,
            Tuple[
                ProduceCallable,
                fastkafka._aiokafka_imports.AIOKafkaProducer,
                Optional[KafkaBrokers],
                Dict[str, Any],
            ],
        ] = {}

        self._producers_list: List[fastkafka._aiokafka_imports.AIOKafkaProducer] = []  # type: ignore

        self.benchmark_results: Dict[str, Dict[str, Any]] = {}

        # background tasks
        self._scheduled_bg_tasks: List[Callable[..., Coroutine[Any, Any, Any]]] = []
        self._bg_task_group_generator: Optional[anyio.abc.TaskGroup] = None
        self._bg_tasks_group: Optional[anyio.abc.TaskGroup] = None

        # todo: use this for errrors
        self._on_error_topic: Optional[str] = None

        self.lifespan = lifespan
        self.lifespan_ctx: Optional[AsyncContextManager[None]] = None

        self._is_started: bool = False
        self._is_shutting_down: bool = False
        self._kafka_consumer_tasks: List[asyncio.Task[Any]] = []
        self._kafka_producer_tasks: List[asyncio.Task[Any]] = []
        self._running_bg_tasks: List[asyncio.Task[Any]] = []
        self.run = False

        # testing functions
        self.AppMocks = None
        self.mocks = None
        self.awaited_mocks = None

    @property
    def is_started(self) -> bool:
        """Property indicating whether the FastKafka object is started.

        The is_started property indicates if the FastKafka object is currently
        in a started state. This implies that all background tasks, producers,
        and consumers have been initiated, and the object is successfully connected
        to the Kafka broker.

        Returns:
            bool: True if the object is started, False otherwise.
        """
        return self._is_started

    def set_kafka_broker(self, kafka_broker_name: str) -> None:
        """
        Sets the Kafka broker to start FastKafka with

        Args:
            kafka_broker_name: The name of the Kafka broker to start FastKafka

        Raises:
            ValueError: If the provided kafka_broker_name is not found in dictionary of kafka_brokers
        """

        if kafka_broker_name not in self._kafka_brokers.brokers:
            raise ValueError(
                f"Given kafka_broker_name '{kafka_broker_name}' is not found in kafka_brokers, available options are {self._kafka_brokers.brokers.keys()}"
            )

        self._kafka_config["bootstrap_servers_id"] = kafka_broker_name

    async def __aenter__(self) -> "FastKafka":
        if self.lifespan is not None:
            self.lifespan_ctx = self.lifespan(self)
            await self.lifespan_ctx.__aenter__()
        await self._start()
        return self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[types.TracebackType],
    ) -> None:
        await self._stop()
        if self.lifespan_ctx is not None:
            await self.lifespan_ctx.__aexit__(exc_type, exc, tb)

    async def _start(self) -> None:
        raise NotImplementedError

    async def _stop(self) -> None:
        raise NotImplementedError

    def consumes(
        self,
        topic: Optional[str] = None,
        decoder: str = "json",
        *,
        prefix: str = "on_",
        brokers: Optional[KafkaBrokers] = None,
        description: Optional[str] = None,
        **kwargs: Dict[str, Any],
    ) -> ConsumeCallable:
        raise NotImplementedError

    def produces(
        self,
        topic: Optional[str] = None,
        encoder: str = "json",
        *,
        prefix: str = "to_",
        brokers: Optional[KafkaBrokers] = None,
        description: Optional[str] = None,
        **kwargs: Dict[str, Any],
    ) -> ProduceCallable:
        raise NotImplementedError

    def benchmark(
        self,
        interval: Union[int, timedelta] = 1,
        *,
        sliding_window_size: Optional[int] = None,
    ) -> Callable[[F], F]:
        raise NotImplementedError

    def run_in_background(
        self,
    ) -> Callable[[], Any]:
        raise NotImplementedError

    def _populate_consumers(
        self,
        is_shutting_down_f: Callable[[], bool],
    ) -> None:
        raise NotImplementedError

    def get_topics(self) -> Iterable[str]:
        raise NotImplementedError

    async def _populate_producers(self) -> None:
        raise NotImplementedError

    async def _populate_bg_tasks(self) -> None:
        raise NotImplementedError

    def create_docs(self) -> None:
        raise NotImplementedError

    def create_mocks(self) -> None:
        raise NotImplementedError

    async def _shutdown_consumers(self) -> None:
        raise NotImplementedError

    async def _shutdown_producers(self) -> None:
        raise NotImplementedError

    async def _shutdown_bg_tasks(self) -> None:
        raise NotImplementedError

# %% ../../nbs/015_FastKafka.ipynb 27
def _get_decoder_fn(decoder: str) -> Callable[[bytes, Type[BaseModel]], Any]:
    """
    Imports and returns decoder function based on input
    """
    if decoder == "json":
        from fastkafka._components.encoder.json import json_decoder

        return json_decoder
    elif decoder == "avro":
        try:
            from fastkafka._components.encoder.avro import avro_decoder
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "Unable to import avro packages. Please install FastKafka using the command 'fastkafka[avro]'"
            )
        return avro_decoder
    else:
        raise ValueError(f"Unknown decoder - {decoder}")

# %% ../../nbs/015_FastKafka.ipynb 29
def _prepare_and_check_brokers(
    app: FastKafka, kafka_brokers: Optional[Union[Dict[str, Any], KafkaBrokers]]
) -> Optional[KafkaBrokers]:
    if kafka_brokers is not None:
        prepared_brokers = _get_kafka_brokers(kafka_brokers)
        if prepared_brokers.brokers.keys() != app._kafka_brokers.brokers.keys():
            raise ValueError(
                f"To override application default brokers, you must define all of the broker options. Default defined: {set(app._kafka_brokers.brokers.keys())}, override defined: {set(prepared_brokers.brokers.keys())}"
            )
        return prepared_brokers
    return None

# %% ../../nbs/015_FastKafka.ipynb 30
def _resolve_key(key: str, dictionary: Dict[str, Any]) -> str:
    i = 0
    resolved_key = f"{key}_{i}"
    while resolved_key in dictionary:
        i += 1
        resolved_key = f"{key}_{i}"
    return resolved_key

# %% ../../nbs/015_FastKafka.ipynb 31
@patch
@delegates(fastkafka._aiokafka_imports.AIOKafkaConsumer)
def consumes(
    self: FastKafka,
    topic: Optional[str] = None,
    decoder: Union[str, Callable[[bytes, Type[BaseModel]], Any]] = "json",
    *,
    executor: Union[str, StreamExecutor, None] = None,
    brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = None,
    prefix: str = "on_",
    description: Optional[str] = None,
    **kwargs: Dict[str, Any],
) -> Callable[[ConsumeCallable], ConsumeCallable]:
    """Decorator registering the callback called when a message is received in a topic.

    This function decorator is also responsible for registering topics for AsyncAPI specificiation and documentation.

    Args:
        topic: Kafka topic that the consumer will subscribe to and execute the
            decorated function when it receives a message from the topic,
            default: None. If the topic is not specified, topic name will be
            inferred from the decorated function name by stripping the defined prefix
        decoder: Decoder to use to decode messages consumed from the topic,
                default: json - By default, it uses json decoder to decode
                bytes to json string and then it creates instance of pydantic
                BaseModel. It also accepts custom decoder function.
        executor: Type of executor to choose for consuming tasks. Avaliable options
                are "SequentialExecutor" and "DynamicTaskExecutor". The default option is
                "SequentialExecutor" which will execute the consuming tasks sequentially.
                If the consuming tasks have high latency it is recommended to use
                "DynamicTaskExecutor" which will wrap the consuming functions into tasks
                and run them in on asyncio loop in background. This comes with a cost of
                increased overhead so use it only in cases when your consume functions have
                high latency such as database queries or some other type of networking.
        prefix: Prefix stripped from the decorated function to define a topic name
                if the topic argument is not passed, default: "on_". If the decorated
                function name is not prefixed with the defined prefix and topic argument
                is not passed, then this method will throw ValueError
        brokers: Optional argument specifying multiple broker clusters for consuming
                messages from different Kafka clusters in FastKafka.
        description: Optional description of the consuming function async docs.
                If not provided, consuming function __doc__ attr will be used.

    Returns:
        A function returning the same function

    Throws:
        ValueError

    """

    def _decorator(
        on_topic: ConsumeCallable,
        topic: Optional[str] = topic,
        decoder: Union[str, Callable[[bytes, Type[BaseModel]], Any]] = decoder,
        executor: Union[str, StreamExecutor, None] = executor,
        brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = brokers,
        description: Optional[str] = description,
        kwargs: Dict[str, Any] = kwargs,
    ) -> ConsumeCallable:
        topic_resolved: str = (
            _get_topic_name(topic_callable=on_topic, prefix=prefix)
            if topic is None
            else topic
        )

        decoder_fn = _get_decoder_fn(decoder) if isinstance(decoder, str) else decoder

        prepared_broker = _prepare_and_check_brokers(self, brokers)
        if prepared_broker is not None:
            self._override_brokers.append(prepared_broker.brokers)  # type: ignore

        if description is not None:
            setattr(on_topic, "description", description)

        self._consumers_store[_resolve_key(topic_resolved, self._consumers_store)] = (
            on_topic,
            decoder_fn,
            executor,
            prepared_broker,
            kwargs,
        )
        setattr(self, on_topic.__name__, on_topic)
        return on_topic

    return _decorator

# %% ../../nbs/015_FastKafka.ipynb 34
def _get_encoder_fn(encoder: str) -> Callable[[BaseModel], bytes]:
    """
    Imports and returns encoder function based on input
    """
    if encoder == "json":
        from fastkafka._components.encoder.json import json_encoder

        return json_encoder
    elif encoder == "avro":
        try:
            from fastkafka._components.encoder.avro import avro_encoder
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "Unable to import avro packages. Please install FastKafka using the command 'fastkafka[avro]'"
            )
        return avro_encoder
    else:
        raise ValueError(f"Unknown encoder - {encoder}")

# %% ../../nbs/015_FastKafka.ipynb 36
@patch
@delegates(fastkafka._aiokafka_imports.AIOKafkaProducer)
def produces(
    self: FastKafka,
    topic: Optional[str] = None,
    encoder: Union[str, Callable[[BaseModel], bytes]] = "json",
    *,
    prefix: str = "to_",
    brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = None,
    description: Optional[str] = None,
    **kwargs: Dict[str, Any],
) -> Callable[[ProduceCallable], ProduceCallable]:
    """Decorator registering the callback called when delivery report for a produced message is received

    This function decorator is also responsible for registering topics for AsyncAPI specificiation and documentation.

    Args:
        topic: Kafka topic that the producer will send returned values from
            the decorated function to, default: None- If the topic is not
            specified, topic name will be inferred from the decorated function
            name by stripping the defined prefix.
        encoder: Encoder to use to encode messages before sending it to topic,
                default: json - By default, it uses json encoder to convert
                pydantic basemodel to json string and then encodes the string to bytes
                using 'utf-8' encoding. It also accepts custom encoder function.
        prefix: Prefix stripped from the decorated function to define a topic
            name if the topic argument is not passed, default: "to_". If the
            decorated function name is not prefixed with the defined prefix
            and topic argument is not passed, then this method will throw ValueError
        brokers: Optional argument specifying multiple broker clusters for consuming
            messages from different Kafka clusters in FastKafka.
        description: Optional description of the producing function async docs.
                If not provided, producing function __doc__ attr will be used.

    Returns:
        A function returning the same function

    Raises:
        ValueError: when needed
    """

    def _decorator(
        to_topic: ProduceCallable,
        topic: Optional[str] = topic,
        brokers: Optional[Union[Dict[str, Any], KafkaBrokers]] = brokers,
        description: Optional[str] = description,
        kwargs: Dict[str, Any] = kwargs,
    ) -> ProduceCallable:
        topic_resolved: str = (
            _get_topic_name(topic_callable=to_topic, prefix=prefix)
            if topic is None
            else topic
        )

        topic_key = _resolve_key(topic_resolved, self._producers_store)

        prepared_broker = _prepare_and_check_brokers(self, brokers)
        if prepared_broker is not None:
            self._override_brokers.append(prepared_broker.brokers)  # type: ignore

        if description is not None:
            setattr(to_topic, "description", description)

        self._producers_store[topic_key] = (
            to_topic,
            None,
            prepared_broker,
            kwargs,
        )
        encoder_fn = _get_encoder_fn(encoder) if isinstance(encoder, str) else encoder
        decorated = producer_decorator(
            self._producers_store,
            to_topic,
            topic_key,
            encoder_fn=encoder_fn,
        )
        setattr(self, to_topic.__name__, decorated)
        return decorated

    return _decorator

# %% ../../nbs/015_FastKafka.ipynb 39
@patch
def get_topics(self: FastKafka) -> Iterable[str]:
    """
    Get all topics for both producing and consuming.

    Returns:
        A set of topics for both producing and consuming.
    """
    produce_topics = set([remove_suffix(topic) for topic in self._producers_store])
    consume_topics = set([remove_suffix(topic) for topic in self._consumers_store])
    return consume_topics.union(produce_topics)

# %% ../../nbs/015_FastKafka.ipynb 41
@patch
def run_in_background(
    self: FastKafka,
) -> Callable[
    [Callable[..., Coroutine[Any, Any, Any]]], Callable[..., Coroutine[Any, Any, Any]]
]:
    """
    Decorator to schedule a task to be run in the background.

    This decorator is used to schedule a task to be run in the background when the app's `_on_startup` event is triggered.

    Returns:
        Callable[None, None]: A decorator function that takes a background task as an input and stores it to be run in the backround.
    """

    def _decorator(
        bg_task: Callable[..., Coroutine[Any, Any, Any]]
    ) -> Callable[..., Coroutine[Any, Any, Any]]:
        """
        Store the background task.

        Args:
            bg_task (Callable[[], None]): The background task to be run asynchronously.

        Returns:
            Callable[[], None]: Original background task.
        """
        logger.info(
            f"run_in_background() : Adding function '{bg_task.__name__}' as background task"
        )
        self._scheduled_bg_tasks.append(bg_task)

        return bg_task

    return _decorator

# %% ../../nbs/015_FastKafka.ipynb 45
@patch
def _populate_consumers(
    self: FastKafka,
    is_shutting_down_f: Callable[[], bool],
) -> None:
    default_config: Dict[str, Any] = filter_using_signature(
        fastkafka._aiokafka_imports.AIOKafkaConsumer, **self._kafka_config
    )

    bootstrap_server = self._kafka_config["bootstrap_servers_id"]

    self._kafka_consumer_tasks = [
        asyncio.create_task(
            aiokafka_consumer_loop(
                topic="_".join(topic.split("_")[:-1]),
                decoder_fn=decoder_fn,
                callback=consumer,
                msg_type=signature(consumer).parameters["msg"].annotation,
                is_shutting_down_f=is_shutting_down_f,
                executor=executor,
                **{
                    **default_config,
                    **override_config,
                    **{
                        "bootstrap_servers": _get_broker_addr_list(
                            kafka_brokers.brokers[bootstrap_server]
                            if kafka_brokers is not None
                            else self._kafka_brokers.brokers[bootstrap_server]
                        )
                    },
                },
            )
        )
        for topic, (
            consumer,
            decoder_fn,
            executor,
            kafka_brokers,
            override_config,
        ) in self._consumers_store.items()
    ]


@patch
async def _shutdown_consumers(
    self: FastKafka,
) -> None:
    if self._kafka_consumer_tasks:
        await asyncio.wait(self._kafka_consumer_tasks)

# %% ../../nbs/015_FastKafka.ipynb 47
# TODO: Add passing of vars
async def _create_producer(  # type: ignore
    *,
    callback: ProduceCallable,
    default_config: Dict[str, Any],
    override_config: Dict[str, Any],
    bootstrap_servers: Union[str, List[str]],
    producers_list: List[fastkafka._aiokafka_imports.AIOKafkaProducer],
) -> fastkafka._aiokafka_imports.AIOKafkaProducer:
    """Creates a producer

    Args:
        callback: A callback function that is called when the producer is ready.
        producer: An existing producer to use.
        default_config: A dictionary of default configuration values.
        override_config: A dictionary of configuration values to override.
        bootstrap_servers: Bootstrap servers to connect the producer to.
        producers_list: A list of producers to add the new producer to.

    Returns:
        A producer.
    """

    config = {
        **filter_using_signature(
            fastkafka._aiokafka_imports.AIOKafkaProducer, **default_config
        ),
        **filter_using_signature(
            fastkafka._aiokafka_imports.AIOKafkaProducer, **override_config
        ),
        **{"bootstrap_servers": bootstrap_servers},
    }

    producer = fastkafka._aiokafka_imports.AIOKafkaProducer(**config)
    logger.info(
        f"_create_producer() : created producer using the config: '{sanitize_kafka_config(**config)}'"
    )

    await producer.start()

    producers_list.append(producer)

    return producer


@patch
async def _populate_producers(self: FastKafka) -> None:
    """Populates the producers for the FastKafka instance.

    Args:
        self: The FastKafka instance.

    Returns:
        None.

    Raises:
        None.
    """
    default_config: Dict[str, Any] = self._kafka_config
    bootstrap_server = default_config["bootstrap_servers_id"]

    self._producers_list = []
    self._producers_store.update(
        {
            topic: (
                callback,
                await _create_producer(
                    callback=callback,
                    default_config=default_config,
                    override_config=override_config,
                    bootstrap_servers=_get_broker_addr_list(
                        kafka_brokers.brokers[bootstrap_server]
                        if kafka_brokers is not None
                        else self._kafka_brokers.brokers[bootstrap_server]
                    ),
                    producers_list=self._producers_list,
                ),
                kafka_brokers,
                override_config,
            )
            for topic, (
                callback,
                _,
                kafka_brokers,
                override_config,
            ) in self._producers_store.items()
        }
    )


@patch
async def _shutdown_producers(self: FastKafka) -> None:
    [await producer.stop() for producer in self._producers_list[::-1]]
    # Remove references to stale producers
    self._producers_list = []
    self._producers_store.update(
        {
            topic: (
                callback,
                None,
                kafka_brokers,
                override_config,
            )
            for topic, (
                callback,
                _,
                kafka_brokers,
                override_config,
            ) in self._producers_store.items()
        }
    )

# %% ../../nbs/015_FastKafka.ipynb 49
@patch
async def _populate_bg_tasks(
    self: FastKafka,
) -> None:
    def _start_bg_task(task: Callable[..., Coroutine[Any, Any, Any]]) -> asyncio.Task:
        logger.info(
            f"_populate_bg_tasks() : Starting background task '{task.__name__}'"
        )
        return asyncio.create_task(task(), name=task.__name__)

    self._running_bg_tasks = [_start_bg_task(task) for task in self._scheduled_bg_tasks]


@patch
async def _shutdown_bg_tasks(
    self: FastKafka,
) -> None:
    for task in self._running_bg_tasks:
        logger.info(
            f"_shutdown_bg_tasks() : Cancelling background task '{task.get_name()}'"
        )
        task.cancel()

    for task in self._running_bg_tasks:
        logger.info(
            f"_shutdown_bg_tasks() : Waiting for background task '{task.get_name()}' to finish"
        )
        try:
            await task
        except asyncio.CancelledError:
            pass
        logger.info(
            f"_shutdown_bg_tasks() : Execution finished for background task '{task.get_name()}'"
        )

# %% ../../nbs/015_FastKafka.ipynb 51
@patch
async def _start(self: FastKafka) -> None:
    def is_shutting_down_f(self: FastKafka = self) -> bool:
        return self._is_shutting_down

    #     self.create_docs()
    await self._populate_producers()
    self._populate_consumers(is_shutting_down_f)
    await self._populate_bg_tasks()

    self._is_started = True


@patch
async def _stop(self: FastKafka) -> None:
    self._is_shutting_down = True

    await self._shutdown_bg_tasks()
    await self._shutdown_consumers()
    await self._shutdown_producers()

    self._is_shutting_down = False
    self._is_started = False

# %% ../../nbs/015_FastKafka.ipynb 57
@patch
def create_docs(self: FastKafka) -> None:
    """
    Create the asyncapi documentation based on the configured consumers and producers.

    This function exports the asyncapi specification based on the configured consumers
    and producers in the FastKafka instance. It generates the asyncapi documentation by
    extracting the topics and callbacks from the consumers and producers.

    Note:
        The asyncapi documentation is saved to the location specified by the `_asyncapi_path`
        attribute of the FastKafka instance.
    """
    (self._asyncapi_path / "docs").mkdir(exist_ok=True, parents=True)
    (self._asyncapi_path / "spec").mkdir(exist_ok=True, parents=True)
    export_async_spec(
        consumers={
            remove_suffix(topic) if topic.endswith("_0") else topic: callback
            for topic, (callback, _, _, _, _) in self._consumers_store.items()
        },
        producers={
            remove_suffix(topic) if topic.endswith("_0") else topic: callback
            for topic, (callback, _, _, _) in self._producers_store.items()
        },
        kafka_brokers=self._kafka_brokers,
        kafka_service_info=self._kafka_service_info,
        asyncapi_path=self._asyncapi_path,
    )

# %% ../../nbs/015_FastKafka.ipynb 61
class AwaitedMock:
    """
    Class representing an awaited mock object.

    Args:
        o: The original object to be wrapped.
    """

    @staticmethod
    def _await_for(f: Callable[..., Any]) -> Callable[..., Any]:
        @delegates(f)
        async def inner(
            *args: Any, f: Callable[..., Any] = f, timeout: int = 60, **kwargs: Any
        ) -> Any:
            """
            Decorator to await the execution of a function.

            Args:
                f: The function to be wrapped.

            Returns:
                The wrapped function.
            """
            if inspect.iscoroutinefunction(f):
                return await asyncio.wait_for(f(*args, **kwargs), timeout=timeout)
            else:
                t0 = datetime.now()
                e: Optional[Exception] = None
                while True:
                    try:
                        return f(*args, **kwargs)
                    except Exception as _e:
                        await asyncio.sleep(1)
                        e = _e

                    if datetime.now() - t0 > timedelta(seconds=timeout):
                        break

                raise e

        return inner

    def __init__(self, o: Any):
        """
        Initializes an instance of AwaitedMock.

        Args:
            o: The original object to be wrapped.
        """
        self._o = o

        for name in o.__dir__():
            if not name.startswith("_"):
                f = getattr(o, name)
                if inspect.ismethod(f):
                    setattr(self, name, self._await_for(f))

# %% ../../nbs/015_FastKafka.ipynb 62
@patch
def create_mocks(self: FastKafka) -> None:
    """Creates self.mocks as a named tuple mapping a new function obtained by calling the original functions and a mock"""
    app_methods = [f for f, _, _, _, _ in self._consumers_store.values()] + [
        f for f, _, _, _ in self._producers_store.values()
    ]
    self.AppMocks = namedtuple(  # type: ignore
        f"{self.__class__.__name__}Mocks", [f.__name__ for f in app_methods]
    )

    self.mocks = self.AppMocks(  # type: ignore
        **{
            f.__name__: AsyncMock() if inspect.iscoroutinefunction(f) else MagicMock()
            for f in app_methods
        }
    )

    self.awaited_mocks = self.AppMocks(  # type: ignore
        **{name: AwaitedMock(mock) for name, mock in self.mocks._asdict().items()}
    )

    def add_mock(
        f: Callable[..., Any], mock: Union[AsyncMock, MagicMock]
    ) -> Callable[..., Any]:
        """Add call to mock when calling function f"""

        @functools.wraps(f)
        async def async_inner(
            *args: Any, f: Callable[..., Any] = f, mock: AsyncMock = mock, **kwargs: Any
        ) -> Any:
            await mock(*args, **kwargs)
            return await f(*args, **kwargs)

        @functools.wraps(f)
        def sync_inner(
            *args: Any, f: Callable[..., Any] = f, mock: MagicMock = mock, **kwargs: Any
        ) -> Any:
            mock(*args, **kwargs)
            return f(*args, **kwargs)

        if inspect.iscoroutinefunction(f):
            return async_inner
        else:
            return sync_inner

    self._consumers_store.update(
        {
            name: (
                add_mock(f, getattr(self.mocks, f.__name__)),
                decoder_fn,
                executor,
                kafka_brokers,
                kwargs,
            )
            for name, (
                f,
                decoder_fn,
                executor,
                kafka_brokers,
                kwargs,
            ) in self._consumers_store.items()
        }
    )

    self._producers_store.update(
        {
            name: (
                add_mock(f, getattr(self.mocks, f.__name__)),
                producer,
                kafka_brokers,
                kwargs,
            )
            for name, (
                f,
                producer,
                kafka_brokers,
                kwargs,
            ) in self._producers_store.items()
        }
    )

# %% ../../nbs/015_FastKafka.ipynb 67
@patch
def benchmark(
    self: FastKafka,
    interval: Union[int, timedelta] = 1,
    *,
    sliding_window_size: Optional[int] = None,
) -> Callable[[Callable[[I], Optional[O]]], Callable[[I], Optional[O]]]:
    """Decorator to benchmark produces/consumes functions

    Args:
        interval: Period to use to calculate throughput. If value is of type int,
            then it will be used as seconds. If value is of type timedelta,
            then it will be used as it is. default: 1 - one second
        sliding_window_size: The size of the sliding window to use to calculate
            average throughput. default: None - By default average throughput is
            not calculated
    """

    def _decorator(func: Callable[[I], Optional[O]]) -> Callable[[I], Optional[O]]:
        func_name = f"{func.__module__}.{func.__qualname__}"

        @wraps(func)
        def wrapper(
            *args: I,
            **kwargs: I,
        ) -> Optional[O]:
            _benchmark(
                interval=interval,
                sliding_window_size=sliding_window_size,
                func_name=func_name,
                benchmark_results=self.benchmark_results,
            )
            return func(*args, **kwargs)

        @wraps(func)
        async def async_wrapper(
            *args: I,
            **kwargs: I,
        ) -> Optional[O]:
            _benchmark(
                interval=interval,
                sliding_window_size=sliding_window_size,
                func_name=func_name,
                benchmark_results=self.benchmark_results,
            )
            return await func(*args, **kwargs)  # type: ignore

        if inspect.iscoroutinefunction(func):
            return async_wrapper  # type: ignore
        else:
            return wrapper

    return _decorator

# %% ../../nbs/015_FastKafka.ipynb 69
@patch
def fastapi_lifespan(
    self: FastKafka, kafka_broker_name: str
) -> Callable[["FastAPI"], AsyncIterator[None]]:
    """
    Method for managing the lifespan of a FastAPI application with a specific Kafka broker.

    Args:
        kafka_broker_name: The name of the Kafka broker to start FastKafka

    Returns:
        Lifespan function to use for initializing FastAPI
    """

    @asynccontextmanager
    async def lifespan(fastapi_app: "FastAPI") -> AsyncIterator[None]:
        self.set_kafka_broker(kafka_broker_name=kafka_broker_name)
        async with self:
            yield

    return lifespan  # type: ignore
