import hashlib
import math

from fastecdsa.curve import P256
from fastecdsa.point import Point
from fastecdsa import util

Q = 31415926535987932384626433832795*P256.G

class EccRNG():
    def __init__(self, Q, seed, curve = P256):
        self.curve = curve
        self.state = int.from_bytes(seed, "big", seed, signed=False )
        self.Q = Q
        self.P = curve.G
        self.tmp = None
        assert Q.curve == curve

    def gen(self):
        new_point = self.state * self.P
        r = new_point.x
        rQ = r * self.P
        random_int_to_return = int(str(bin((rQ).x))[16:], 2)
        self.state = (r*self.Q).x
        self.lsb = str(bin((rQ).x))
        self.rQ = rQ
        return random_int_to_return.to_bytes(math.ceil(random_int_to_return.bit_length()/8), byteorder='big')

class LazarusKeyScheduler():
    def __init__(self, password: bytes, nonce: bytes):
        self.ecc = EccRNG()
        self.hashState = hashlib.sha3_256(password + nonce + password).digest()
        self.eccState = self.ecc.gen()
        for i in range(64):
            self._stateStep()

    def genToken(self):
        self._stateStep()
        return self.state + self.eccState

    def getKey(self, length: int):
        return hashlib.shake_256(self.genToken()).digest(length)

    def _stateStep(self):
        self.eccState = self.ecc.gen()
        self.hashState = hashlib.sha3_256(self.hashState + self.eccState).digest()

class OmegaChain():
    def __init__(self, key: bytes, primer: bytes):
        self.blockSize = 16
        self.prevBlock = primer
        self.key = key

    def encrypt(self, data: bytes):
        blocks = self._dataToBlocks(data)
        enc = bytes()
        for block in blocks:
            enc += self.encBlock(block)

    def _encBlock(self, plainBlock: bytes):
        chainedBlock = self._xor(self.prevBlock, plainBlock)
        encBlock = self._xor(chainedBlock, self.key)
        self.prevBlock = encBlock
        return encBlock

    def _xor(self, a: bytes, b: bytes):
        if len(a)!=self.blockSize or len(b)!=self.blockSize:
            raise Exception("Cannot xor") #TODO
        return bytes([a[i]^b[i] for i in range(len(a))])

    def _dataToBlocks(self, data: bytes):
        padded = self._padData(data, blockSize = self.blockSize)
        return [data[b*self.blockSize:self.blockSize] for b in range(int(len(padded)/self.blockSize))]

    def _padData(self, data: bytes, blockSize: int = None):
        if not blockSize:
            blockSize = self.blockSize
        bytesMissing = (blockSize-len(data)%blockSize)
        if bytesMissing == blockSize:
            return data
        else:
            return data + bytes(bytesMissing)

class Lazarus():
    def encrypt(password: bytes, plaintext: bytes, nonce: bytes = None):
        if not nonce:
            #TODO: Random iv
            nonce = b'bla'
        scheduler = LazarusKeyScheduler(password, nonce)
        omegaKey = scheduler.genToken()
        primer = scheduler.getKey(16)
        aesKey = scheduler.getKey(16)
        aesIV = scheduler.getKey(16)
        hmacInnerKey = scheduler.getKey(16)
        hmacOuterKey = scheduler.getKey(16)

        alpha = Lazarus._encryptAES(plaintext, aesKey, aesIV)
        hmac = Lazarus._genHMAC(plaintext, hmacInnerKey, hmacOuterKey)
        psi = alpha + hmac
        omega = OmegaChain(omegaKey, primer)
        return omega.encrypt(psi)

    def omegaChain(prev: bytes, scheduler: LazarusKeyScheduler):
        pass

    def eccPerm(self, key: bytes):
        pass

    def primePerm(self, key: bytes):
        pass