# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/001_ConsumerLoop.ipynb.

# %% auto 0
__all__ = ['logger', 'process_msgs', 'process_message_callback', 'aiokafka_consumer_loop']

# %% ../../nbs/001_ConsumerLoop.ipynb 1
import asyncio
from asyncio import iscoroutinefunction  # do not use the version from inspect
from datetime import datetime, timedelta
from os import environ
from typing import *

import anyio
from anyio.streams.memory import MemoryObjectReceiveStream
import asyncer
from aiokafka import AIOKafkaConsumer
from aiokafka.structs import ConsumerRecord, TopicPartition
from pydantic import BaseModel, Field, HttpUrl, NonNegativeInt

from .logger import get_logger

# %% ../../nbs/001_ConsumerLoop.ipynb 6
logger = get_logger(__name__)

# %% ../../nbs/001_ConsumerLoop.ipynb 10
async def process_msgs(  # type: ignore
    *,
    msgs: Dict[TopicPartition, List[ConsumerRecord]],
    callbacks: Dict[str, Callable[[BaseModel], Union[None, Awaitable[None]]]],
    msg_types: Dict[str, Type[BaseModel]],
    process_f: Callable[
        [Tuple[Callable[[BaseModel], Awaitable[None]], BaseModel]], Awaitable[None]
    ],
) -> None:
    """For each messages **msg** in **msgs**, calls process_f with callbacks[topic] and **msgs**.

    Params:
        msgs: a dictionary mapping topic partition to a list of messages, returned by `AIOKafkaConsumer.getmany`.
        callbacks: a dictionary mapping topics into a callback functions.
        msg_types: a dictionary mapping topics into a message type of a message.
        process_f: a stream processing function registrated by `anyio.create_memory_object_stream`

    Todo:
        remove it :)
    """
    for topic_partition, topic_msgs in msgs.items():
        topic = topic_partition.topic
        msg_type = msg_types[topic]
        try:
            decoded_msgs = [
                msg_type.parse_raw(msg.value.decode("utf-8")) for msg in topic_msgs
            ]
            for msg in decoded_msgs:
                callback_raw = callbacks[topic]
                if not iscoroutinefunction(callback_raw):
                    c: Callable[[BaseModel], None] = callback_raw  # type: ignore
                    callback: Callable[[BaseModel], Awaitable[None]] = asyncer.asyncify(
                        c
                    )
                else:
                    callback = callback_raw

                async def safe_callback(
                    msg: BaseModel,
                    callback: Callable[[BaseModel], Awaitable[None]] = callback,
                ) -> None:
                    try:
                        #                         logger.debug(f"process_msgs(): awaiting '{callback}({msg})'")
                        await callback(msg)
                    except Exception as e:
                        logger.warning(
                            f"process_msgs(): exception caugth {e.__repr__()} while awaiting '{callback}({msg})'"
                        )

                await process_f((safe_callback, msg))
        except Exception as e:
            logger.warning(
                f"process_msgs(): Unexpected exception '{e.__repr__()}' caught and ignored for topic='{topic_partition.topic}', partition='{topic_partition.partition}' and messages: {topic_msgs}"
            )

# %% ../../nbs/001_ConsumerLoop.ipynb 17
async def process_message_callback(
    receive_stream: MemoryObjectReceiveStream[Any],
) -> None:
    async with receive_stream:
        async for callback, msg in receive_stream:
            await callback(msg)


async def _aiokafka_consumer_loop(  # type: ignore
    consumer: AIOKafkaConsumer,
    *,
    callbacks: Dict[str, Callable[[BaseModel], Union[None, Awaitable[None]]]],
    timeout_ms: int = 100,
    max_buffer_size: int = 10_000,
    msg_types: Dict[str, Type[BaseModel]],
    is_shutting_down_f: Callable[[], bool],
) -> None:
    """Write docs

    Todo: add batch size if needed
    """
    send_stream, receive_stream = anyio.create_memory_object_stream(
        max_buffer_size=max_buffer_size
    )
    async with anyio.create_task_group() as tg:
        tg.start_soon(process_message_callback, receive_stream)
        async with send_stream:
            while not is_shutting_down_f():
                msgs = await consumer.getmany(timeout_ms=timeout_ms)
                try:
                    await process_msgs(
                        msgs=msgs,
                        callbacks=callbacks,
                        msg_types=msg_types,
                        process_f=send_stream.send,
                    )
                except Exception as e:
                    logger.warning(
                        f"_aiokafka_consumer_loop(): Unexpected exception '{e}' caught and ignored for messages: {msgs}"
                    )

# %% ../../nbs/001_ConsumerLoop.ipynb 19
async def aiokafka_consumer_loop(  # type: ignore
    topics: List[str],
    *,
    bootstrap_servers: str,
    auto_offset_reset: str,
    max_poll_records: int = 1_000,
    timeout_ms: int = 100,
    max_buffer_size: int = 10_000,
    callbacks: Dict[str, Callable[[BaseModel], Union[None, Awaitable[None]]]],
    msg_types: Dict[str, Type[BaseModel]],
    is_shutting_down_f: Callable[[], bool],
    **kwargs,
) -> None:
    """todo: write docs"""
    logger.info(f"aiokafka_consumer_loop() starting..")
    consumer = AIOKafkaConsumer(
        bootstrap_servers=bootstrap_servers,
        auto_offset_reset=auto_offset_reset,
        max_poll_records=max_poll_records,
    )
    logger.info("aiokafka_consumer_loop(): Consumer created.")

    await consumer.start()
    logger.info("aiokafka_consumer_loop(): Consumer started.")
    consumer.subscribe(topics)
    logger.info("aiokafka_consumer_loop(): Consumer subscribed.")

    try:
        await _aiokafka_consumer_loop(
            consumer=consumer,
            max_buffer_size=max_buffer_size,
            timeout_ms=timeout_ms,
            callbacks=callbacks,
            msg_types=msg_types,
            is_shutting_down_f=is_shutting_down_f,
        )
    finally:
        await consumer.stop()
        logger.info(f"aiokafka_consumer_loop(): Consumer stopped.")
        logger.info(f"aiokafka_consumer_loop() finished.")
