from typing import List

import kikimr.public.sdk.python.persqueue.grpc_pq_streaming_api as pqlib
from kikimr.public.sdk.python.persqueue.auth import TVMCredentialsProvider
from kikimr.public.sdk.python.persqueue.errors import SessionFailureResult
from ticket_parser2 import (
    TvmApiClientSettings,
    TvmClient
)

from ..settings import logger


class LogBroker:
    def __init__(self,
                 host: str,
                 port: int,
                 topics: List[str],
                 consumer: str,
                 tvm_id: int,
                 tvm_secret: str,
                 dsts_tvm_id: int,
                 wait_timeout: int,
                 max_time_lag_ms: int):
        self.__host = host
        self.__port = port
        self.__topics = topics
        self.__consumer = consumer
        self.__tvm_id = tvm_id
        self.__tvm_secret = tvm_secret
        self.__dsts_tvm_id = dsts_tvm_id
        self.__wait_timeout = wait_timeout
        self.__max_time_lag_ms = max_time_lag_ms

    def _get_streaming_api(self) -> pqlib.PQStreamingAPI:
        api = pqlib.PQStreamingAPI(self.__host, self.__port)
        logger.info('Starting PQLib')
        api_start_future = api.start()
        result = api_start_future.result(timeout=self.__wait_timeout)
        logger.info(f'API started with result: {result}')
        return api

    def _get_credentials(self) -> TVMCredentialsProvider:
        tvm_settings = TvmApiClientSettings(self_client_id=self.__tvm_id,
                                            self_secret=self.__tvm_secret,
                                            dsts={'pq': self.__dsts_tvm_id})
        tvm_client = TvmClient(tvm_settings)
        return TVMCredentialsProvider(tvm_client, destination_alias='pq')

    def _get_consumer(self) -> pqlib.PQStreamingConsumer:
        api = self._get_streaming_api()
        credentials_provider = self._get_credentials()
        configurator = pqlib.ConsumerConfigurator(topics=self.__topics,
                                                  client_id=self.__consumer,
                                                  use_client_locks=True,
                                                  max_time_lag_ms=self.__max_time_lag_ms)
        consumer = api.create_consumer(consumer_configurator=configurator,
                                       credentials_provider=credentials_provider)

        logger.info('Starting consumer')
        start_future = consumer.start()
        start_result = start_future.result(timeout=self.__wait_timeout)
        if not isinstance(start_result, SessionFailureResult):
            if start_result.HasField("init"):
                logger.info(f'Consumer start result was: {start_result}')
            else:
                raise RuntimeError(f'Bad consumer start result from server: {start_result}.')
        else:
            raise RuntimeError(f'Error occurred on start of consumer: {start_result}.')
        logger.info('Consumer started')

        return consumer

    def read(self):
        consumer = self._get_consumer()
        logger.info('Reading data')
        while True:
            try:
                result = consumer.next_event().result()
            except TimeoutError:
                raise RuntimeError('Failed to get any messages from topic. Are your sure there is any data to read?')

            if result.type == pqlib.ConsumerMessageType.MSG_DATA:
                for batch in result.message.data.message_batch:
                    for message in batch.message:
                        yield message
                consumer.commit(result.message.data.cookie)
            elif result.type == pqlib.ConsumerMessageType.MSG_LOCK:
                result.ready_to_read()
                logger.info(f'Got partition assignment: topic {result.message.lock.topic}, partition {result.message.lock.partition}')
            elif result.type == pqlib.ConsumerMessageType.MSG_RELEASE:
                logger.info(f'Partition revoked. Topic {result.message.release.topic}, partition {result.message.release.partition}')
            elif result.type == pqlib.ConsumerMessageType.MSG_ERROR:
                logger.info('Received error, need to recreate consumer')
                raise RuntimeError('Need to recreate consumer becaise MSG_ERROR received')
