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

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.domain.models.auto_create_topic_configuration import AutoCreateTopicConfiguration
from buz.kafka.domain.models.consumer_initial_offset_position import ConsumerInitialOffsetPosition
from buz.kafka.domain.models.kafka_connection_config import KafkaConnectionConfig
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,
        health_check_port: Optional[int],
        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.__health_check_port = health_check_port
        self.__consumer_initial_offset_position = consumer_initial_offset_position
        self.__max_records_retrieved_per_poll = 1
        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 configure_health_check_server(self, health_check_port: int) -> web.TCPSite:
        self._logger.info(f"Starting health check server on port {health_check_port}")
        app = web.Application()
        app.router.add_get("/health", lambda request: self.__health_check())
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, "localhost", health_check_port)
        await site.start()
        return site

    async def run(self) -> None:
        start_time = datetime.now()
        await self.__generate_kafka_consumers()
        health_check_server: Optional[web.TCPSite] = None

        if self.__health_check_port is not None:
            health_check_server = await self.configure_health_check_server(self.__health_check_port)

        self.__initial_coroutines_created_elapsed_time = datetime.now() - start_time

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

        start_consumption_time = datetime.now()
        worker_errors = await self.__run_worker()
        self.__events_processed_elapsed_time = datetime.now() - start_consumption_time

        if health_check_server is not None:
            await health_check_server.stop()

        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
            raise WorkerExecutionException(
                "The worker was closed by an unexpected exception"
            ) from consume_events_exception or polling_task_exception

    async def __run_worker(self) -> tuple[Optional[Exception], Optional[Exception]]:
        consume_events_task = create_task(self.__consume_events_task())
        polling_task = create_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, task: Task) -> Optional[Exception]:
        try:
            await task
            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.__generate_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 __generate_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.__executor_per_consumer_mapper[kafka_consumer] = executor

            self.__queue_per_consumer_mapper[kafka_consumer] = InMemoryMultiqueueRepository()

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

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

    async def __polling_task(self) -> None:
        self._logger.info("Initializing subscribers")
        try:
            polling_task_per_consumer = [
                create_task(self.__polling_consuming_tasks(consumer))
                for consumer, subscriber in self.__queue_per_consumer_mapper.items()
            ]

            await gather(*polling_task_per_consumer)

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

    async def __polling_consuming_tasks(self, consumer: AIOKafkaConsumer) -> None:
        queue = self.__queue_per_consumer_mapper[consumer]

        try:
            self._logger.info(
                f"initializing consumer group: '{consumer.get_consumer_group()}' subscribed to the topics: '{consumer.get_topics()}'"
            )
            await consumer.init()
            self._logger.info(f"initialized '{consumer.get_consumer_group()}'")
        except Exception:
            self._logger.exception(
                f"Unexpected error during Kafka subscriber '{consumer.get_consumer_group()}' initialization. Skipping it: {traceback.format_exc()}"
            )

        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

            async with self.__polling_tasks_semaphore:
                kafka_poll_records = await consumer.poll(
                    number_of_messages_to_poll=self.__max_records_retrieved_per_poll,
                )

                for kafka_poll_record in kafka_poll_records:
                    queue.push(
                        key=TopicPartition(
                            topic=kafka_poll_record.topic,
                            partition=kafka_poll_record.partition,
                        ),
                        record=kafka_poll_record,
                    )

            if len(kafka_poll_records) == 0:
                await sleep(self.__seconds_between_polls_if_there_are_no_new_tasks)

    async def __consume_events_task(self) -> None:
        self._logger.info("Initializing consuming task")
        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[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():
            if await self.__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

                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

    async def __all_queues_are_empty(self) -> bool:
        return all([queue.is_totally_empty() for queue in self.__queue_per_consumer_mapper.values()])

    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.__queue_per_consumer_mapper.keys():
            await kafka_consumer.stop()

    async def __health_check(self) -> web.Response:
        health_information = {
            "subscribers": [subscriber.fqn() for subscriber in self.__subscribers],
            "number_of_subscribers": len(self.__subscribers),
            "event_processed": self.__events_processed,
        }

        return web.Response(text=json.dumps(health_information), content_type="application/json")

    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}")
