from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import traceback

from thrift.protocol.THeaderProtocol import THeaderProtocol
from thrift.server.TServer import TServer, TConnectionContext
from thrift.transport.THeaderTransport import THeaderTransport
from thrift.transport.TTransport import TMemoryBuffer

from thrift.server.CppServerWrapper import CppServerWrapper, CppContextData, \
    SSLPolicy, SSLVerifyPeerEnum, SSLVersion, CallbackWrapper, CallTimestamps

from concurrent.futures import Future
from functools import partial

import time

# Default sampling rate for expensive sampling operations, such as histogram
# counters.
kDefaultSampleRate = 32


class TCppConnectionContext(TConnectionContext):
    def __init__(self, context_data):
        self.context_data = context_data

    def getClientPrincipal(self):
        return self.context_data.getClientIdentity()

    def getClientPrincipalUser(self):
        principal = self.getClientPrincipal()
        if not principal:
            return None
        user, match, domain = principal.partition('@')
        if match:
            return user
        return None

    def getPeerName(self):
        return self.context_data.getPeerAddress()

    def getSockName(self):
        return self.context_data.getLocalAddress()

class _ProcessorAdapter(object):
    CONTEXT_DATA = CppContextData
    CALLBACK_WRAPPER = CallbackWrapper

    def __init__(self, processor):
        self.processor = processor
        self.observer = None
        self.sampleRate = kDefaultSampleRate
        self.sampleCount = 0

    def setObserver(self, observer):
        self.observer = observer

    def _shouldSample(self):
        self.sampleCount = (self.sampleCount + 1) % self.sampleRate
        return self.sampleCount == 0

    # TODO mhorowitz: add read headers here, so they can be added to
    # the constructed header buffer.  Also add endpoint addrs to the
    # context
    def call_processor(self, input, headers, client_type, protocol_type,
                       context_data, callback):
        try:
            # The input string has already had the header removed, but
            # the python processor will expect it to be there.  In
            # order to reconstitute the message with headers, we use
            # the THeaderProtocol object to write into a memory
            # buffer, then pass that buffer to the python processor.

            should_sample = self._shouldSample()

            timestamps = CallTimestamps()
            timestamps.processBegin = 0
            timestamps.processEnd = 0
            if self.observer and should_sample:
                timestamps.processBegin = int(time.time() * 10**6)