import abc
import asyncio
import functools
import grpc
import logging

from typing import TypeVar, Callable, Any, Optional

from cached_property import cached_property

import ydb.public.api.protos.ydb_operation_pb2 as kikimr_operations
import ydb.public.api.protos.ydb_status_codes_pb2 as kikimr_status_codes

from logbroker.public.api.protos.common_pb2 import Path

from saas.library.python.logbroker.common import LogbrokerEndpoint
from saas.library.python.token_store import PersistentTokenStore

TService = TypeVar('TService')
TServiceClass = TypeVar('TServiceClass')


class GRPCClient(metaclass=abc.ABCMeta):
    _MAX_MESSAGE_SIZE = 64 * 10 ** 6
    _CONNECT_OPTIONS = [
        ('grpc.max_receive_message_length', _MAX_MESSAGE_SIZE),
        ('grpc.max_send_message_length', _MAX_MESSAGE_SIZE),
        ('grpc.primary_user_agent', 'CMClient')
    ]
    _TIMEOUT = 5

    @classmethod
    @functools.lru_cache()
    def _get_grpc_channel(cls, endpoint: str) -> grpc.Channel:
        channel = grpc.insecure_channel(endpoint, options=cls._CONNECT_OPTIONS)
        grpc.channel_ready_future(channel).result(timeout=cls._TIMEOUT)
        return channel


class LogbrokerClient(GRPCClient, metaclass=abc.ABCMeta):
    _CM_LOGBROKER_ENDPOINT = LogbrokerEndpoint('cm.logbroker.yandex.net', 1111)

    _PROTO_MAPPING = {}

    _REQUEST_CHECK_DELAY = 1
    _REQUEST_CHECK_MAX_ATTEMPTS = 10

    def __init__(self, token: Optional[str] = None, endpoint: LogbrokerEndpoint = _CM_LOGBROKER_ENDPOINT) -> None:
        if not token:
            token = PersistentTokenStore.get_token_from_store_env_or_file('logbroker')

        self._token: str = token.strip()
        self._endpoint: LogbrokerEndpoint = endpoint

    @classmethod
    async def _wrap_async(cls, sync_func: Callable) -> Any:
        """
        Run a synchronous function in executor.
        WARNING! The function should have its own timeout, otherwise the process will not stop until complete.

        :param sync_func: A synchronous function to be executed
        :return: A result of the function call
        """
        return await asyncio.get_event_loop().run_in_executor(None, sync_func)

    @property
    @abc.abstractmethod
    def _service_stub_cls(self) -> TServiceClass:
        raise NotImplementedError

    @cached_property
    def _service_stub(self) -> TService:
        channel = self._get_grpc_channel(str(self._endpoint))
        return self._service_stub_cls(channel)

    async def __get_response_data(self, service: TService, response: Any, attempt: int = 1) -> Any:
        if attempt > self._REQUEST_CHECK_MAX_ATTEMPTS:
            raise LogbrokerAPIException(
                f'Unable to get response data in {self._REQUEST_CHECK_MAX_ATTEMPTS} attempt(s): {response}'
            )

        if response.operation.status == kikimr_status_codes.StatusIds.SUCCESS and response.operation.ready is True:
            return response.operation.result.value
        elif (
            response.operation.status == kikimr_status_codes.StatusIds.STATUS_CODE_UNSPECIFIED
            and
            response.operation.ready is False
        ):
            await asyncio.sleep(self._REQUEST_CHECK_DELAY)

            request = kikimr_operations.GetOperationRequest(id=response.operation.id)
            response = await self._wrap_async(lambda: service.GetOperation(request, timeout=self._TIMEOUT))

            return await self.__get_response_data(service, response, attempt=attempt + 1)
        else:
            logging.debug(response.operation.status)
            logging.debug(response.operation.ready)

            raise LogbrokerAPIException(f'API ERROR: {response}')

    async def _make_request(self, handle_name, metadata=None, **kwargs) -> Any:
        handle = getattr(self._service_stub, handle_name)
        request_class = self._PROTO_MAPPING[handle_name]

        if 'path' in kwargs.keys() and not isinstance(kwargs['path'], Path):
            kwargs['path'] = Path(path=kwargs['path'])

        request = request_class(
            operation_params=kikimr_operations.OperationParams(
                operation_mode=kikimr_operations.OperationParams.OperationMode.SYNC
            ),
            **kwargs
        )
        response = await self._wrap_async(lambda: handle(request, metadata=metadata, timeout=self._TIMEOUT))
        return await self.__get_response_data(self._service_stub, response)


class LogbrokerAPIException(Exception):
    ...
