import abc
import hmac
import hashlib
import math
class HKDF(object):
    __metaclass__ = abc.ABCMeta
    HASH_OUTPUT_SIZE  = 32

    @staticmethod
    def createFor(messageVersion):
        from . import hkdfv2, hkdfv3
        if messageVersion == 2:
            return hkdfv2.HKDFv2()
        elif messageVersion == 3:
            return hkdfv3.HKDFv3()
        else:
            raise  AssertionError("Unknown version: %s " % messageVersion)

    def deriveSecrets(self, inputKeyMaterial, info, outputLength, salt = None):
        salt = salt or bytearray(self.__class__.HASH_OUTPUT_SIZE)
        prk = self.extract(salt, inputKeyMaterial)
        return self.expand(prk, info, outputLength)

    def extract(self, salt, inputKeyMaterial):
        mac = hmac.new(bytes(salt), digestmod=hashlib.sha256)
        mac.update(bytes(inputKeyMaterial))
        return mac.digest()

    def expand(self, prk, info, outputSize):

        iterations = int(math.ceil(float(outputSize)/ float(self.__class__.HASH_OUTPUT_SIZE)))
        mixin = bytearray()
        results = bytearray()
        remainingBytes = outputSize

        for i in range(self.getIterationStartOffset(), iterations + self.getIterationStartOffset()):
            mac = hmac.new(prk, digestmod=hashlib.sha256)
            mac.update(bytes(mixin))
            if info is not None:
                mac.update(bytes(info))
            updateChr = chr(i % 256)
            mac.update(updateChr.encode())

            stepResult = mac.digest()
            stepSize = min(remainingBytes, len(stepResult))
            results.extend(stepResult[:stepSize])
            mixin = stepResult
            remainingBytes -= stepSize

        return bytes(results)

    @abc.abstractmethod
    def getIterationStartOffset(self):
        return 0
