import logging
import random
import threading
import time
from datetime import datetime
from azure.data.tables import TableClient, TableTransactionError, TableServiceClient
from azure.core.credentials import AzureSasCredential
from common import Options
from common.schedule import Queue, QueueEvent

logger = logging.getLogger(__name__)
MAX_DATE = (datetime(2100, 1, 1)-datetime(1970,1,1)).total_seconds()

class ModelPerformanceLogHandler(logging.Handler):

    def __init__(self, **options):
        super(ModelPerformanceLogHandler, self).__init__()
        self.options = Options(**options)
        self.export_interval = self.options.export_interval
        self.max_batch_size = self.options.max_batch_size
        self.level = logging.DEBUG
        self.addFilter(SamplingFilter(self.options.logging_sampling_rate))
        self._queue = Queue(capacity=self.options.queue_capacity)
        self._worker = Worker(self._queue, self)
        service = TableServiceClient(
            endpoint=self.options.endpoint, credential=AzureSasCredential(self.options.sas_token))
        self._table_client = service.get_table_client(
            table_name=self.options.table_name)
        self._worker.start()

    def _export(self, batch, event=None):  # pragma: NO COVER
        try:
            if batch:
                envelopes = [self.log_record_to_envelope(x) for x in batch]
                try:
                    self._table_client.submit_transaction(envelopes)
                except TableTransactionError as e:
                    logger.exception(
                        "There was an error with model performace trace.")
        finally:
            if event:
                event.set()

    def _export_fallback(self, batch, event=None):  # pragma: NO COVER
        try:
            if batch:
                envelopes = [self.log_record_to_envelope(x) for x in batch]
                try:
                    for envelope in envelopes: 
                        self._table_client.create_entity(entity=envelope[1])
                except TableTransactionError as e:
                    logger.exception(
                        "There was an error with model performace trace.")
        finally:
            if event:
                event.set()

    # Close is automatically called as part of logging shutdown
    def close(self, timeout=None):
        if not timeout:
            timeout = self.options.grace_period
        if self._worker:
            self._worker.stop(timeout)
        super(ModelPerformanceLogHandler, self).close()

    def createLock(self):
        self.lock = None

    def emit(self, record):
        self._queue.put(record, block=False)

    def log_record_to_envelope(self, record):
        properties = {}
        properties.update(record.model_performance)
        # AttributeError: 'dict' object has no attribute 'PartitionKey'
        properties['PartitionKey'] = record.model_signature
        properties['message'] = record.message
        total_ms = MAX_DATE - record.created
        properties['RowKey'] = "%018d" % total_ms
        return ("create", properties)

    def flush(self, timeout=None):
        if self._queue.is_empty():
            return

        if not self._worker.is_alive():
            logger.warning("Can't flush %s, worker thread is dead. "
                           "Any pending messages will be lost.", self)
            return

        self._queue.flush(timeout=timeout)


class SamplingFilter(logging.Filter):

    def __init__(self, probability=1.0):
        super(SamplingFilter, self).__init__()
        self.probability = probability

    def filter(self, record):
        if (hasattr(record, 'model_signature') and hasattr(record, 'model_performance') and isinstance(record.model_performance, dict)):
            return random.random() < self.probability
        else:
            return False


class Worker(threading.Thread):
    daemon = True

    def __init__(self, src, dst):
        self._src = src
        self._dst = dst
        self._stopping = False
        super(Worker, self).__init__(
            name='{} Worker'.format(type(dst).__name__)
        )

    def run(self):
        src = self._src
        dst = self._dst
        while True:
            batch = src.gets(dst.max_batch_size, dst.export_interval)

            if batch and batch[-1] is src.EXIT_EVENT:
                dst._export_fallback(batch[:-1])
                break

            if batch and isinstance(batch[-1], QueueEvent):
                try:
                    dst._export(batch[:-1], event=batch[-1])
                except Exception:
                    logger.exception('Unhandled exception from exporter.')
                if batch[-1] is src.EXIT_EVENT:
                    break
                continue  # pragma: NO COVER
            try:
                dst._export(batch)
            except Exception:
                logger.exception('Unhandled exception from exporter.')

    def stop(self, timeout=None):  # pragma: NO COVER
        start_time = time.time()
        wait_time = timeout
        if self.is_alive() and not self._stopping:
            self._stopping = True
            self._src.put(self._src.EXIT_EVENT, block=True, timeout=wait_time)
            elapsed_time = time.time() - start_time
            wait_time = timeout and max(timeout - elapsed_time, 0)
        if self._src.EXIT_EVENT.wait(timeout=wait_time):
            return time.time() - start_time  # time taken to stop
