# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/012_ProducerManager.ipynb.

# %% auto 0
__all__ = ['logger', 'AIOKafkaProducerManager']

# %% ../../nbs/012_ProducerManager.ipynb 1
import asyncio
from contextlib import asynccontextmanager, contextmanager
from typing import *

import anyio
from aiokafka import AIOKafkaProducer
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream

from .logger import get_logger

# %% ../../nbs/012_ProducerManager.ipynb 4
logger = get_logger(__name__)

# %% ../../nbs/012_ProducerManager.ipynb 6
@asynccontextmanager
async def _aiokafka_producer_manager(  # type: ignore # Argument 1 to "_aiokafka_producer_manager" becomes "Any" due to an unfollowed import  [no-any-unimported]
    producer: AIOKafkaProducer,
    *,
    max_buffer_size: int = 10_000,
) -> AsyncGenerator[MemoryObjectSendStream[Any], None]:
    """Write docs

    Todo: add batch size if needed

    """

    logger.info("_aiokafka_producer_manager(): Starting...")

    async def send_message(receive_stream: MemoryObjectReceiveStream) -> Any:
        async with receive_stream:
            async for topic, msg in receive_stream:
                fut = await producer.send(topic, msg)
                msg = await fut

    send_stream, receive_stream = anyio.create_memory_object_stream(
        max_buffer_size=max_buffer_size
    )

    logger.info("_aiokafka_producer_manager(): Starting send_stream")
    asyncio.create_task(send_message(receive_stream))
    async with send_stream:
        yield send_stream
        logger.info("_aiokafka_producer_manager(): Exiting send_stream")

    logger.info("_aiokafka_producer_manager(): Finished.")

# %% ../../nbs/012_ProducerManager.ipynb 9
class AIOKafkaProducerManager:
    def __init__(self, producer: AIOKafkaProducer, *, max_buffer_size: int = 1_000):  # type: ignore
        self.producer = producer
        self.max_buffer_size = max_buffer_size

    async def start(self) -> None:
        logger.info("AIOKafkaProducerManager.start(): Entering...")
        await self.producer.start()
        self.producer_manager_generator = _aiokafka_producer_manager(self.producer)
        self.send_stream = await self.producer_manager_generator.__aenter__()
        logger.info("AIOKafkaProducerManager.start(): Finished.")

    async def stop(self) -> None:
        # todo: try to flush messages before you exit
        logger.info("AIOKafkaProducerManager.stop(): Entering...")
        await self.producer_manager_generator.__aexit__(None, None, None)
        logger.info("AIOKafkaProducerManager.stop(): Stoping producer...")
        await self.producer.stop()
        logger.info("AIOKafkaProducerManager.stop(): Finished")

    def send(self, topic: str, msg: bytes) -> None:
        self.send_stream.send_nowait((topic, msg))
