from .gtfo_table_app import GtfoTableApp, TableTransaction
from .gtfo_app import NoMessageError
from nubium_utils.confluent_utils.confluent_runtime_vars import env_vars
from confluent_kafka import TopicPartition
import logging
from copy import deepcopy
from json import loads, dumps
from datetime import datetime


LOGGER = logging.getLogger(__name__)


class AutoBatchTableTransaction(TableTransaction):
    def __init__(self, producer, consumer, app_changelog_topic, app_tables, metrics_manager=None, message=None, auto_consume=True, timeout=5, parent_app=None):
        self._commit_offsets = {}
        self.message_count = 0
        super().__init__(producer, consumer, app_changelog_topic=app_changelog_topic, app_tables=app_tables, metrics_manager=metrics_manager, message=message, auto_consume=auto_consume, timeout=timeout, parent_app=parent_app)
        self._pending_table_writes = {p: {} for p in self.app_tables}
        self._pending_table_offset_increase = {p: 0 for p in self.app_tables}

    def _refresh_transaction_fields(self):
        super()._refresh_transaction_fields()
        self._pending_table_writes = {p: {} for p in self.app_tables}
        self._pending_table_offset_increase = {p: 0 for p in self.app_tables}
        self._commit_offsets = {}
        self.message_count = 0

    def read_table_entry(self):
        pending_update = self._pending_table_writes.get(self.partition(), {}).get(self.key())
        if pending_update:
            return deepcopy(loads(pending_update))
        return super().read_table_entry()

    def update_table_entry(self, value,  mark_changelog_updated=False):
        super().update_table_entry(value, mark_changelog_updated=mark_changelog_updated)
        self._update_pending_table_writes()

    def delete_table_entry(self):
        super().delete_table_entry()
        self._update_pending_table_writes()

    def _update_pending_table_writes(self):
        self._pending_table_writes[self.partition()][self.key()] = self._pending_table_write
        self._pending_table_offset_increase[self.partition()] += 1

    def _table_write(self):
        # TODO: have this share objects with table_app to consolidate logic
        for p, msgs in self._pending_table_writes.items():
            writes = {}
            for key, msg in msgs.items():
                if msg == '-DELETED-':
                    LOGGER.debug('Finalizing table entry delete...')
                    self.app_tables[p].delete(key)
                else:
                    writes[key] = msg
            if writes:
                table_offset = str(self._table_offset(partition=p) + 1)
                writes['offset'] = table_offset  # the +1 is because the transaction causes an extra offset at the end (per trans + partition)
                LOGGER.debug(f'Finalizing table entry batch write for p{p}')
                self.app_tables[p].write_batch(writes)
        self._pending_table_write = None
        self._pending_table_writes = {}
        self._pending_table_offset_increase = {}

    def _table_offset(self, partition=None):
        if partition is None:
            partition = self.partition()
        return super()._table_offset(partition) + self._pending_table_offset_increase.get(partition, 0)

    def consume(self):
        super().consume()
        self._commit_offsets[self.partition()] = {"topic": self.topic(), "partition": self.partition(), "offset": self.offset()}
        self.message_count += 1

    def commit(self, mark_committed=True):
        # TODO: Confirm there aren't commits from old partition assignments (not sure how that's handled...it may not rebalance
        #  while doing a transaction?)
        offsets_to_commit = [TopicPartition(msg['topic'], msg['partition'], msg['offset'] + 1) for msg in
                             self._commit_offsets.values()]
        if offsets_to_commit:
            LOGGER.debug(f'Offsets to commit: {offsets_to_commit}')
            self._init_transaction()
        if self._active_transaction:
            LOGGER.info(f'Committing {self.message_count} messages...')
            self.producer.send_offsets_to_transaction(offsets_to_commit, self.consumer.consumer_group_metadata())
            self.producer.commit_transaction(30)
            self.producer.poll(0)  # TODO: experiment with flushing? I think it does it by default
            LOGGER.info('All partition transactions committed!')
            self._committed = mark_committed
            self._commit_offsets = {}
            self.message_count = 0
            self._retry_offsets = {}
            if self._committed:
                self._active_transaction = False
        self._table_write()


class GtfoAutoBatchTableApp(GtfoTableApp):
    def __init__(self, app_function, consume_topic, produce_topic_schema_dict=None, transaction_type=AutoBatchTableTransaction,
                 app_function_arglist=None, metrics_manager=None, schema_registry=None, cluster_name=None, consumer=None, producer=None,
                 time_elapse_max_seconds=None, consume_max_count=None):

        self._time_elapse_start = None

        if time_elapse_max_seconds is None:
            time_elapse_max_seconds = int(env_vars()['NU_CONSUMER_DEFAULT_BATCH_CONSUME_MAX_TIME_SECONDS'])
        if consume_max_count is None:
            consume_max_count = int(env_vars()['NU_CONSUMER_DEFAULT_BATCH_CONSUME_MAX_COUNT'])

        self._time_elapse_max_seconds = time_elapse_max_seconds
        self._consume_max_count = consume_max_count

        super().__init__(
            app_function, consume_topic, produce_topic_schema_dict, transaction_type=transaction_type,
            app_function_arglist=app_function_arglist, metrics_manager=metrics_manager, schema_registry=schema_registry, cluster_name=cluster_name, consumer=consumer, producer=producer)

    def prepare_transaction(self, *args, **kwargs):
        return super().consume(*args, auto_consume=False, **kwargs)

    def _max_consume_time_continue(self):
        if self._time_elapse_max_seconds:
            seconds_elapsed = datetime.now().timestamp() - self._time_elapse_start
            return seconds_elapsed < self._time_elapse_max_seconds
        return True

    def _max_consume_count_continue(self, consume_multiplier=1):
        if self._consume_max_count:
            return self.transaction.message_count < (self._consume_max_count * consume_multiplier)
        return True

    def _keep_consuming(self, consume_multiplier=1):
        return not self._shutdown and self._max_consume_count_continue(consume_multiplier=consume_multiplier) and self._max_consume_time_continue()

    def _check_pending_recovery_status(self):
        LOGGER.debug(f'Current pending recovery partition statuses before refreshing: {self._pending_table_recoveries}')
        recovery_status = {p: {'table_offset': self.transaction._table_offset(partition=p), 'highwater': self._pending_table_recoveries[p]['watermarks'][1]} for p in self._pending_table_recoveries.keys()}
        LOGGER.info(f'CURRENT RECOVERY STATUS:\n{dumps(recovery_status, indent=4)}')
        for p in recovery_status:
            if recovery_status[p]['highwater'] - (recovery_status[p]['table_offset'] + 2) <= 0:
                LOGGER.info('Offset delta for is within an expected tolerable range of the changelog highwater; it is safe to assume transactional markers are the remaining delta.')
                self._set_table_offset_to_latest(p)

    def _table_recovery_loop(self, checks=2, consume_multiplier=10):
        while self._pending_table_recoveries and checks:
            LOGGER.info(f'Consuming from changelog partitions: {list(self._active_table_changelog_partitions.keys())}')
            self.prepare_transaction()
            if self._time_elapse_max_seconds:
                self._time_elapse_start = datetime.now().timestamp()
            LOGGER.info(f'Processing up to {self._consume_max_count * consume_multiplier} messages for up to {self._time_elapse_max_seconds} seconds!')
            try:
                while self._keep_consuming(consume_multiplier=consume_multiplier) and self._pending_table_recoveries and checks:
                    self.transaction.consume()
                    self.transaction._recover_table_via_changelog()
                    # NOTE: no commit here since its just writing to the table
            except NoMessageError:
                checks -= 1
                LOGGER.debug(f'No changelog messages, checks remaining: {checks}')
            self.transaction._table_write()
            self._check_pending_recovery_status()
        if not checks:
            for p in list(self._pending_table_recoveries.keys()):  # list stops dict iter size change exception
                self._set_table_offset_to_latest(p)

    def _app_run(self, *args, **kwargs):
        LOGGER.info(f'Consuming from partitions: {list(self._active_primary_partitions.keys())}')
        self.prepare_transaction(*args, **kwargs)
        if self._time_elapse_max_seconds:
            self._time_elapse_start = datetime.now().timestamp()
        LOGGER.info(
            f'Processing up to {self._consume_max_count} messages for up to {self._time_elapse_max_seconds} seconds!')
        try:
            while self._keep_consuming():
                self.transaction.consume()
                self.app_function(self.transaction, *self.app_function_arglist)
        except NoMessageError:
            self.producer.poll(0)
            LOGGER.info('No messages!')
            pass
        if any(self.transaction._pending_table_offset_increase.values()) or self.transaction._commit_offsets:
            self.transaction.commit()
