# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/013_ProducerDecorator.ipynb.

# %% auto 0
__all__ = ['BaseSubmodel', 'ProduceReturnTypes', 'ProduceCallable', 'KafkaEvent', 'get_loop', 'release_callback',
           'produce_single', 'send_batch', 'produce_batch', 'producer_decorator']

# %% ../../nbs/013_ProducerDecorator.ipynb 1
import random
import asyncio
import functools
import json
import time
from asyncio import iscoroutinefunction  # do not use the version from inspect
from collections import namedtuple
from dataclasses import dataclass
from typing import *

import nest_asyncio
from aiokafka import AIOKafkaProducer
from aiokafka.producer.message_accumulator import BatchBuilder
from pydantic import BaseModel

from .meta import export

# %% ../../nbs/013_ProducerDecorator.ipynb 3
BaseSubmodel = TypeVar("BaseSubmodel", bound=Union[List[BaseModel], BaseModel])
BaseSubmodel


@dataclass
@export("fastkafka")
class KafkaEvent(Generic[BaseSubmodel]):
    """
    A generic class for representing Kafka events. Based on BaseSubmodel, bound to pydantic.BaseModel

    Attributes:
        message (BaseSubmodel): The message contained in the Kafka event, can be of type pydantic.BaseModel.
        key (bytes, optional): The optional key used to identify the Kafka event.
    """

    message: BaseSubmodel
    key: Optional[bytes] = None

# %% ../../nbs/013_ProducerDecorator.ipynb 5
ProduceReturnTypes = Union[
    BaseModel, KafkaEvent[BaseModel], List[BaseModel], KafkaEvent[List[BaseModel]]
]

ProduceCallable = Union[
    Callable[..., ProduceReturnTypes], Callable[..., Awaitable[ProduceReturnTypes]]
]

# %% ../../nbs/013_ProducerDecorator.ipynb 8
def _wrap_in_event(
    message: Union[BaseModel, List[BaseModel], KafkaEvent]
) -> KafkaEvent:
    return message if type(message) == KafkaEvent else KafkaEvent(message)

# %% ../../nbs/013_ProducerDecorator.ipynb 11
def get_loop() -> asyncio.AbstractEventLoop:
    try:
        loop: asyncio.AbstractEventLoop = asyncio.get_event_loop()
    except RuntimeError as e:
        loop = asyncio.new_event_loop()

    if loop.is_running():
        nest_asyncio.apply(loop)

    return loop

# %% ../../nbs/013_ProducerDecorator.ipynb 13
def release_callback(fut: asyncio.Future) -> None:
    pass

# %% ../../nbs/013_ProducerDecorator.ipynb 14
async def produce_single(  # type: ignore
    producer: AIOKafkaProducer,
    topic: str,
    encoder_fn: Callable[[BaseModel], bytes],
    wrapped_val: KafkaEvent[BaseModel],
) -> ProduceReturnTypes:
    fut = await producer.send(
        topic, encoder_fn(wrapped_val.message), key=wrapped_val.key
    )
    fut.add_done_callback(release_callback)

# %% ../../nbs/013_ProducerDecorator.ipynb 16
async def send_batch(  # type: ignore
    producer: AIOKafkaProducer, topic: str, batch: BatchBuilder, key: Optional[bytes]
) -> None:
    partitions = await producer.partitions_for(topic)
    if key == None:
        partition = random.choice(tuple(partitions))  # nosec
    else:
        partition = producer._partition(topic, None, None, None, key, None)
    await producer.send_batch(batch, topic, partition=partition)


async def produce_batch(  # type: ignore
    producer: AIOKafkaProducer,
    topic: str,
    encoder_fn: Callable[[BaseModel], bytes],
    wrapped_val: KafkaEvent[List[BaseModel]],
) -> ProduceReturnTypes:
    batch = producer.create_batch()

    for message in wrapped_val.message:
        metadata = batch.append(
            key=wrapped_val.key,
            value=encoder_fn(message),
            timestamp=int(time.time() * 1000),
        )
        if metadata == None:
            # send batch
            await send_batch(producer, topic, batch, wrapped_val.key)
            # create new batch
            batch = producer.create_batch()
            batch.append(
                key=None, value=encoder_fn(message), timestamp=int(time.time() * 1000)
            )

    await send_batch(producer, topic, batch, wrapped_val.key)

# %% ../../nbs/013_ProducerDecorator.ipynb 18
def producer_decorator(
    producer_store: Dict[str, Any],
    func: ProduceCallable,
    topic: str,
    encoder_fn: Callable[[BaseModel], bytes],
) -> ProduceCallable:
    """todo: write documentation"""

    loop = get_loop()

    @functools.wraps(func)
    async def _produce_async(
        *args: List[Any],
        topic: str = topic,
        encoder_fn: Callable[[BaseModel], bytes] = encoder_fn,
        producer_store: Dict[str, Any] = producer_store,
        f: Callable[..., Awaitable[ProduceReturnTypes]] = func,  # type: ignore
        **kwargs: Any
    ) -> ProduceReturnTypes:
        return_val = await f(*args, **kwargs)
        wrapped_val = _wrap_in_event(return_val)
        _, producer, _ = producer_store[topic]

        if isinstance(wrapped_val.message, list):
            await produce_batch(producer, topic, encoder_fn, wrapped_val)
        else:
            await produce_single(producer, topic, encoder_fn, wrapped_val)
        return return_val

    @functools.wraps(func)
    def _produce_sync(
        *args: List[Any],
        topic: str = topic,
        encoder_fn: Callable[[BaseModel], bytes] = encoder_fn,
        producer_store: Dict[str, Any] = producer_store,
        f: Callable[..., ProduceReturnTypes] = func,  # type: ignore
        loop: asyncio.AbstractEventLoop = loop,
        **kwargs: Any
    ) -> ProduceReturnTypes:
        return_val = f(*args, **kwargs)
        wrapped_val = _wrap_in_event(return_val)
        _, producer, _ = producer_store[topic]
        if isinstance(wrapped_val.message, list):
            loop.run_until_complete(
                produce_batch(producer, topic, encoder_fn, wrapped_val)
            )
        else:
            loop.run_until_complete(
                produce_single(producer, topic, encoder_fn, wrapped_val)
            )
        return return_val

    return _produce_async if iscoroutinefunction(func) else _produce_sync
