import os
import kikimr.public.sdk.python.persqueue.grpc_pq_streaming_api as pqlib
import kikimr.public.sdk.python.persqueue.auth as auth
from tvmauth import BlackboxTvmId as BlackboxClientId
from tvm2 import TVM2
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.log_reader.lib import log_parser

logger = get_basic_logger(__name__)


class LogBrokerConsumer(object):
    def __init__(self, endpoint, topic, consumer_id):
        self._api_port = 2135
        self._endpoint = endpoint
        self._topic = topic
        self._consumer_id = consumer_id
        self._parser = log_parser.ScoreApiLogParser()
        self._consumer = None
        self._cookies = None
        self._reset()

    def _get_cred_provider(self, use_tvm):
        if use_tvm:
            tvm_client_id = os.environ.get('TVM_CLIENT_ID')
            tvm_secret = os.environ.get('TVM_SECRET')
            assert tvm_client_id, '[LogReader Consumer] Please provide TVM_CLIENT_ID'
            assert tvm_secret, '[LogReader Consumer] Please provide TVM_SECRET'
            tvm_client = TVM2(
                client_id=tvm_client_id,
                secret=tvm_secret,
                blackbox_client=BlackboxClientId.Prod,
                allowed_clients=(),
                destinations=(2001059,),
            )
            cred_provider = auth.TVMCredentialsProvider(tvm_client, destination_client_id=2001059)
        else:
            token = os.environ.get('LOGBROKER_TOKEN')
            assert token, '[LogClient] Please provide LOGBROKER_TOKEN'
            cred_provider = auth.OAuthTokenCredentialsProvider(token)
        return cred_provider

    def _create_consumer(self, use_tvm=True):
        logger.info(' Creating consumer')
        api = pqlib.PQStreamingAPI(self._endpoint, self._api_port)
        api_start_future = api.start()
        result = api_start_future.result(timeout=10)
        logger.info(' Api started with result: {}'.format(result))

        configurator = pqlib.ConsumerConfigurator(self._topic, self._consumer_id)
        cred_provider = self._get_cred_provider(use_tvm)
        consumer = api.create_consumer(configurator, credentials_provider=cred_provider)

        logger.info(' Start consumer')
        start_future = consumer.start()
        start_result = start_future.result(timeout=10)

        try:
            start_result.HasField('init')
            logger.info(' Consumer start result was: {}'.format(start_result))
        except Exception as ex:
            raise RuntimeError(' Consumer failed to start with error {} {}'.format(start_result, ex))
        logger.info(' Consumer started')
        return consumer

    def _process_single_batch(self, consumer_message):
        data = []
        for batch in consumer_message.data.message_batch:
            for message in batch.message:
                data.extend(self._parser.parse_raw_line(message.data))
                logger.info(
                    ' Received message from offset {} with seq_no {}'.format(
                        message.offset, message.meta.seq_no))
        return data

    def read(self, total_messages_expected=100, timeout=5):
        if self._consumer is not None or self._cookies is not None:
            logger.warn('WARNING! Have you call the `commit` function before new `read`?')
        self._consumer = self._create_consumer()
        data = []
        self._cookies = []

        logger.info(' Read data')
        while total_messages_expected > 0:
            result = self._consumer.next_event().result(timeout=timeout)
            assert result.type == pqlib.ConsumerMessageType.MSG_DATA
            extracted_messages = self._process_single_batch(result.message)
            data.extend(extracted_messages)
            logger.info(' Read {} messages'.format(len(extracted_messages)))

            total_messages_expected -= len(extracted_messages)
            self._cookies.append(result.message.data.cookie)

        logger.info(' Reads done, totaly read {} messages'.format(len(data)))
        self._consumer.reads_done()
        return data

    def commit(self, timeout=5):
        logger.info(' Committing cookie now')
        assert self._consumer, 'Call `read` before `commit`, consumer must exist here'
        assert self._cookies is not None, 'Call `read` before `commit`, cookies must be not None'
        if not self._cookies:
            logger.info(' No cookies to commit!')
            self._consumer.stop()
            return
        self._consumer.commit(self._cookies)
        last_commited_cookie = None
        logger.info(' Wait for ack on commits')
        while last_commited_cookie != self._cookies[-1]:
            result = self._consumer.next_event().result(timeout=timeout)
            if result.type == pqlib.ConsumerMessageType.MSG_DATA:
                continue
            assert result.type == pqlib.ConsumerMessageType.MSG_COMMIT
            if last_commited_cookie is not None:
                assert result.message.commit.cookie[-1] > last_commited_cookie
            last_commited_cookie = result.message.commit.cookie[-1]
            logger.info(' All cookies up to {} acks'.format(last_commited_cookie))
        logger.info(' Stop consumer')
        self._consumer.stop()
        self._reset()

    def _reset(self):
        self._consumer = None
        self._cookies = None
