from abc import abstractmethod
import traceback
from asyncio import Lock, gather, Semaphore, Event as AsyncIOEvent, sleep
from datetime import timedelta, datetime
from itertools import cycle
from logging import Logger
from typing import AsyncIterator, Coroutine, Optional, Sequence, Type, TypeVar

from aiokafka import TopicPartition
from aiokafka.coordinator.assignors.abstract import AbstractPartitionAssignor
from aiokafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor

from buz.event import Event
from buz.event.async_consumer import AsyncConsumer
from buz.event.exceptions.worker_execution_exception import WorkerExecutionException
from buz.event.infrastructure.buz_kafka.consume_strategy.consume_strategy import KafkaConsumeStrategy
from buz.event.infrastructure.buz_kafka.kafka_event_subscriber_executor import KafkaEventSubscriberExecutor
from buz.event.infrastructure.models.consuming_task import ConsumingTask
from buz.event.meta_subscriber import MetaSubscriber
from buz.kafka import (
    KafkaConnectionConfig,
    ConsumerInitialOffsetPosition,
)
from buz.kafka.domain.models.auto_create_topic_configuration import AutoCreateTopicConfiguration
from buz.kafka.domain.models.kafka_poll_record import KafkaPollRecord
from buz.kafka.domain.services.kafka_admin_client import KafkaAdminClient
from buz.kafka.infrastructure.aiokafka.aiokafka_consumer import AIOKafkaConsumer
from buz.queue.in_memory.in_memory_multiqueue_repository import InMemoryMultiqueueRepository
from buz.queue.multiqueue_repository import MultiqueueRepository


T = TypeVar("T", bound=Event)


class BaseBuzAIOKafkaAsyncConsumer(AsyncConsumer):
    __FALLBACK_PARTITION_ASSIGNORS = (RoundRobinPartitionAssignor,)

    def __init__(
        self,
        *,
        connection_config: KafkaConnectionConfig,
        kafka_admin_client: Optional[KafkaAdminClient],
        consume_strategy: KafkaConsumeStrategy,
        max_queue_size: int,
        max_records_retrieved_per_poll: int,
        kafka_partition_assignors: tuple[Type[AbstractPartitionAssignor], ...] = (),
        subscribers: Sequence[MetaSubscriber],
        logger: Logger,
        consumer_initial_offset_position: ConsumerInitialOffsetPosition,
        auto_create_topic_configuration: Optional[AutoCreateTopicConfiguration] = None,
        seconds_between_executions_if_there_are_no_tasks_in_the_queue: int = 1,
        seconds_between_polls_if_there_are_tasks_in_the_queue: int = 1,
        seconds_between_polls_if_there_are_no_new_tasks: int = 1,
        max_number_of_concurrent_polling_tasks: int = 20,
    ):
        self.__connection_config = connection_config
        self.__consume_strategy = consume_strategy
        self.__kafka_partition_assignors = kafka_partition_assignors
        self.__subscribers = subscribers
        self._logger = logger
        self.__consumer_initial_offset_position = consumer_initial_offset_position
        self.__max_records_retrieved_per_poll = 1
        self.__subscriber_per_consumer_mapper: dict[AIOKafkaConsumer, MetaSubscriber] = {}
        self.__executor_per_consumer_mapper: dict[AIOKafkaConsumer, KafkaEventSubscriberExecutor] = {}
        self.__queue_per_consumer_mapper: dict[
            AIOKafkaConsumer, MultiqueueRepository[TopicPartition, KafkaPollRecord]
        ] = {}
        self.__max_records_retrieved_per_poll = max_records_retrieved_per_poll
        self.__max_queue_size = max_queue_size
        self.__should_stop = AsyncIOEvent()
        self.__start_kafka_consumers_elapsed_time: Optional[timedelta] = None
        self.__initial_coroutines_created_elapsed_time: Optional[timedelta] = None
        self.__events_processed: int = 0
        self.__events_processed_elapsed_time: timedelta = timedelta()
        self.__kafka_admin_client = kafka_admin_client
        self.__auto_create_topic_configuration = auto_create_topic_configuration
        self.__seconds_between_executions_if_there_are_no_tasks_in_the_queue = (
            seconds_between_executions_if_there_are_no_tasks_in_the_queue
        )
        self.__seconds_between_polls_if_there_are_tasks_in_the_queue = (
            seconds_between_polls_if_there_are_tasks_in_the_queue
        )
        self.__seconds_between_polls_if_there_are_no_new_tasks = seconds_between_polls_if_there_are_no_new_tasks
        self.__polling_tasks_semaphore = Semaphore(max_number_of_concurrent_polling_tasks)

        self.__task_execution_mutex = Lock()

    async def run(self) -> None:
        start_time = datetime.now()
        await self.__generate_kafka_consumers()

        if len(self.__subscriber_per_consumer_mapper) == 0:
            self._logger.error("There are no valid subscribers to execute, finalizing consumer")
            return

        self.__create_queue_repository_per_consumer()
        self.__initial_coroutines_created_elapsed_time = datetime.now() - start_time

        start_consumption_time = datetime.now()
        self._logger.info("Starting to consume events")
        worker_errors = await self.__run_worker()
        self.__events_processed_elapsed_time = datetime.now() - start_consumption_time

        await self.__handle_graceful_stop(worker_errors)

    async def __handle_graceful_stop(self, worker_errors: tuple[Optional[Exception], Optional[Exception]]) -> None:
        self._logger.info("Stopping kafka consumers...")
        await self.__manage_kafka_consumers_stopping()
        self._logger.info("All kafka consumers stopped")

        self.__print_statistics()

        if self.__exceptions_are_thrown(worker_errors):
            consume_events_exception, polling_task_exception = worker_errors
            if consume_events_exception:
                self._logger.error(consume_events_exception)
            if polling_task_exception:
                self._logger.error(polling_task_exception)

            raise WorkerExecutionException("The worker was closed by an unexpected exception")

    async def __run_worker(self) -> tuple[Optional[Exception], Optional[Exception]]:
        consume_events_task = self.__consume_events_task()
        polling_task = self.__polling_task()

        try:
            await gather(consume_events_task, polling_task)
            return (None, None)
        except Exception:
            self.__should_stop.set()
            consume_events_exception = await self.__await_exception(consume_events_task)
            polling_task_exception = await self.__await_exception(polling_task)
            return (consume_events_exception, polling_task_exception)

    async def __await_exception(self, future: Coroutine) -> Optional[Exception]:
        try:
            await future
            return None
        except Exception as exception:
            return exception

    def __exceptions_are_thrown(self, worker_errors: tuple[Optional[Exception], Optional[Exception]]) -> bool:
        return any([error is not None for error in worker_errors])

    async def __generate_kafka_consumers(self):
        start_time = datetime.now()
        tasks = [self.__initialize_kafka_consumer_for_subscriber(subscriber) for subscriber in self.__subscribers]
        await gather(*tasks)
        self.__start_kafka_consumers_elapsed_time = datetime.now() - start_time

    async def __initialize_kafka_consumer_for_subscriber(self, subscriber: MetaSubscriber) -> None:
        try:
            executor = await self._create_kafka_consumer_executor(subscriber)
            topics = self.__consume_strategy.get_topics(subscriber)
            kafka_consumer = AIOKafkaConsumer(
                consumer_group=self.__consume_strategy.get_subscription_group(subscriber),
                topics=topics,
                connection_config=self.__connection_config,
                initial_offset_position=self.__consumer_initial_offset_position,
                partition_assignors=self.__kafka_partition_assignors + self.__FALLBACK_PARTITION_ASSIGNORS,
                logger=self._logger,
                kafka_admin_client=self.__kafka_admin_client,
                auto_create_topic_configuration=self.__auto_create_topic_configuration,
                on_partition_revoked=self.__on_partition_revoked,
            )

            self.__subscriber_per_consumer_mapper[kafka_consumer] = subscriber

            self.__executor_per_consumer_mapper[kafka_consumer] = executor

            await kafka_consumer.init()
        except Exception:
            self._logger.exception(
                f"Unexpected error during Kafka subscriber '{subscriber.fqn()}' initialization. Skipping it: {traceback.format_exc()}"
            )

    @abstractmethod
    async def _create_kafka_consumer_executor(self, subscriber: MetaSubscriber) -> KafkaEventSubscriberExecutor:
        pass

    def __create_queue_repository_per_consumer(self) -> None:
        for kafka_consumer in self.__subscriber_per_consumer_mapper.keys():
            self.__queue_per_consumer_mapper[kafka_consumer] = InMemoryMultiqueueRepository()

    async def __polling_task(self) -> None:
        try:
            while not self.__should_stop.is_set():
                total_size = sum([queue.get_total_size() for queue in self.__queue_per_consumer_mapper.values()])
                if total_size >= self.__max_queue_size:
                    await sleep(self.__seconds_between_polls_if_there_are_tasks_in_the_queue)
                    continue

                raw_consuming_tasks = await gather(
                    *[
                        self.__polling_consuming_tasks(kafka_consumer=consumer)
                        for consumer, subscriber in self.__subscriber_per_consumer_mapper.items()
                    ]
                )

                poll_results: list[ConsumingTask] = [
                    consuming_task for consuming_tasks in raw_consuming_tasks for consuming_task in consuming_tasks
                ]
                if len(poll_results) == 0:
                    await sleep(self.__seconds_between_polls_if_there_are_no_new_tasks)

                for poll_result in poll_results:
                    queue = self.__queue_per_consumer_mapper[poll_result.consumer]
                    queue.push(
                        key=TopicPartition(
                            topic=poll_result.kafka_poll_record.topic, partition=poll_result.kafka_poll_record.partition
                        ),
                        record=poll_result.kafka_poll_record,
                    )

        except Exception:
            self._logger.error(f"Polling task failed with exception: {traceback.format_exc()}")
            self.__should_stop.set()

        return

    async def __polling_consuming_tasks(self, kafka_consumer: AIOKafkaConsumer) -> list[ConsumingTask]:
        async with self.__polling_tasks_semaphore:
            results = await kafka_consumer.poll(
                number_of_messages_to_poll=self.__max_records_retrieved_per_poll,
            )
        return [ConsumingTask(kafka_consumer, result) for result in results]

    async def __consume_events_task(self) -> None:
        blocked_tasks_iterator = self.generate_blocked_consuming_tasks_iterator()

        async for consuming_task in blocked_tasks_iterator:
            consumer = consuming_task.consumer
            kafka_poll_record = consuming_task.kafka_poll_record
            executor = self.__executor_per_consumer_mapper[consuming_task.consumer]

            await executor.consume(kafka_poll_record=kafka_poll_record)

            await consumer.commit_poll_record(kafka_poll_record)

            self.__events_processed += 1

    # This iterator return a blocked task, that will be blocked for other process (like rebalancing), until the next task will be requested
    async def generate_blocked_consuming_tasks_iterator(self) -> AsyncIterator[ConsumingTask]:
        consumer_queues_cyclic_iterator = cycle(self.__queue_per_consumer_mapper.items())
        last_consumer, _ = next(consumer_queues_cyclic_iterator)

        while not self.__should_stop.is_set():
            all_queues_are_empty = all(
                [queue.is_totally_empty() for queue in self.__queue_per_consumer_mapper.values()]
            )

            if all_queues_are_empty:
                await sleep(self.__seconds_between_executions_if_there_are_no_tasks_in_the_queue)
                continue

            async with self.__task_execution_mutex:
                consumer: Optional[AIOKafkaConsumer] = None
                kafka_poll_record: Optional[KafkaPollRecord] = None

                while consumer != last_consumer:
                    consumer, queue = next(consumer_queues_cyclic_iterator)
                    kafka_poll_record = queue.pop()

                    if kafka_poll_record is not None:
                        yield ConsumingTask(consumer, kafka_poll_record)
                        last_consumer = consumer
                        break

        return

    async def __on_partition_revoked(self, consumer: AIOKafkaConsumer, topics_partitions: set[TopicPartition]) -> None:
        async with self.__task_execution_mutex:
            for topic_partition in topics_partitions:
                self.__queue_per_consumer_mapper[consumer].clear(topic_partition)

    def request_stop(self) -> None:
        self.__should_stop.set()
        self._logger.info("Worker stop requested. Waiting for finalize the current task")

    async def __manage_kafka_consumers_stopping(self) -> None:
        for kafka_consumer in self.__subscriber_per_consumer_mapper.keys():
            await kafka_consumer.stop()

    def __print_statistics(self) -> None:
        self._logger.info("Number of subscribers: %d", len(self.__subscribers))
        self._logger.info(f"Start kafka consumers elapsed time: {self.__start_kafka_consumers_elapsed_time}")
        self._logger.info(f"Initial coroutines created elapsed time: {self.__initial_coroutines_created_elapsed_time}")
        self._logger.info(f"Events processed: {self.__events_processed}")
        self._logger.info(f"Events processed elapsed time: {self.__events_processed_elapsed_time}")
