import math
from typing import List, Union, Any

from logbroker.public.api.protos.common_pb2 import ExecuteModifyCommandsResult
from logbroker.public.api.grpc.config_manager_pb2_grpc import ConfigurationManagerServiceStub
from logbroker.public.api.protos.config_manager_pb2 import (
    DescribePathRequest, ListDirectoryRequest, ListYtDeliveriesRequest, ExecuteModifyCommandsRequest,
    DescribePathResult, DescribeAccountResult, DescribeDirectoryResult, DescribeTopicResult, DescribeConsumerResult,
    ListDirectoryResult
)

from saas.library.python.logbroker.common import get_one_of
from saas.library.python.logbroker.configuration.directory import LogbrokerDirectory
from saas.library.python.logbroker.configuration.entities import LogbrokerConsumer
from saas.library.python.logbroker.configuration.topic import LogbrokerTopic
from saas.library.python.logbroker.client import LogbrokerClient, TServiceClass


class ConfigurationClient(LogbrokerClient):

    _PROTO_MAPPING = {  # see logbroker/public/api/grpc/config_manager.proto
        'DescribePath': DescribePathRequest,
        'ListDirectory': ListDirectoryRequest,
        'ListYtDeliveries': ListYtDeliveriesRequest,
        'ExecuteModifyCommands': ExecuteModifyCommandsRequest,
    }

    _ENTITY_TYPE_MAPPING = {
        'topic': LogbrokerTopic,
        'directory': LogbrokerDirectory,
        'consumer': LogbrokerConsumer,
    }

    _DESCRIBE_RESPONSE_MAPPING = {
        'account': DescribeAccountResult,
        'directory': DescribeDirectoryResult,
        'topic': DescribeTopicResult,
        'consumer': DescribeConsumerResult,
    }

    @property
    def _service_stub_cls(self) -> TServiceClass:
        return ConfigurationManagerServiceStub

    async def _make_request(self, handle_name, **kwargs) -> Any:
        return await super()._make_request(handle_name, token=self._token, **kwargs)

    async def list_directory(self, path) -> List[Union[LogbrokerTopic, LogbrokerConsumer, LogbrokerDirectory]]:
        children = []
        raw_response = await self._make_request('ListDirectory', path=path)
        parsed_result = ListDirectoryResult()
        parsed_result.ParseFromString(raw_response)
        for child in parsed_result.children:
            entity_type = child.WhichOneof('entry')
            children.append(self._ENTITY_TYPE_MAPPING[entity_type](self, getattr(child, entity_type)))
        return children

    async def list_topics(self, path, recurse=False) -> List[LogbrokerTopic]:
        entities = await self.list_directory(path)
        topics = [entity for entity in entities if isinstance(entity, LogbrokerTopic)]
        if recurse:
            directories = [entity for entity in entities if isinstance(entity, LogbrokerDirectory)]
            for child in directories:
                child_topics = await self.list_topics(child.path, recurse=True)
                topics.extend(child_topics)
        return topics

    async def describe_path(self, path) -> Union[DescribeTopicResult, DescribeDirectoryResult, DescribeAccountResult]:
        raw_response = await self._make_request('DescribePath', path=path)
        response_message = DescribePathResult()
        response_message.ParseFromString(raw_response)
        result = get_one_of(response_message, 'result')
        return result

    async def execute_modify_commands(self, commands) -> List[ExecuteModifyCommandsResult]:
        def split_batches(arr, batch_size=100):
            for i in range(int(math.ceil(len(commands)/100))):
                slice_begin = i*batch_size
                yield arr[slice_begin:slice_begin + batch_size]

        batches = split_batches(commands)
        results = []
        for batch in batches:
            raw_response = await self._make_request('ExecuteModifyCommands', actions=batch)
            parsed_result = ExecuteModifyCommandsResult()
            parsed_result.ParseFromString(raw_response)
            results.append(parsed_result)
        return results
