# Concept: Create 8 PKs where each represent a bloodtype. Let 7 of them be created by OGen and 1 of them by KeyGen.
# The one represents our bloodtype. Bob will then encrypt 8 values using these PKs, where each value repredents
# A truth value, thus either true or false, s.t. each cipher is an entry in the bloodtype comptability matrix.

from secrets import SystemRandom
import time
from math import pow
import numpy as np
from crypto.week1 import BloodType, convert_from_string_to_enum, blood_cell_compatibility_lookup

# We can't encrypt 0, so we have to index from 1
convert_bloodtype_to_index = {
    BloodType.O_NEGATIVE: 1,
    BloodType.O_POSITIVE: 2,
    BloodType.A_NEGATIVE: 3,
    BloodType.A_POSITIVE: 4,
    BloodType.B_NEGATIVE: 5,
    BloodType.B_POSITIVE: 6,
    BloodType.AB_NEGATIVE: 7,
    BloodType.AB_POSITIVE: 8,
}


class ElGamal:
    def __init__(self, g, q, p):
        self.gen_ = g
        self.order = q
        self.p = p
        self.pk = None
        self.sk = None

    def gen_key(self):
        key = SystemRandom().randint(1, self.order)
        while np.gcd(q, key) != 1:
            key = SystemRandom().randint(1, self.order)
        return key

    def gen(self, sk):
        h = (self.gen_**sk) % self.order
        self.sk = sk
        self.pk = (self.gen_, h)
        return self.pk

    def enc(self, m, pk):
        # sample random r \in Zq
        r = SystemRandom().randint(1, q)
        g, h = pk

        s = (h**r) % q
        p = (g**r) % q
        c = s * m
        return c, p

    def dec(self, c):
        c1, c2 = c
        # c, p, key, q
        h = (c2**self.sk) % q
        m = c1 / h
        return m

    def ogen(self):
        # Here, q = 2p+1, thus we actually need to use the p here, instead of
        # self.order, but as we do not know p yet, .e we
        # TODO: Use p instead of self.order, s.t. self.order = 2p+1
        s = SystemRandom().randint(1, self.order)
        h = s**2 % self.order
        return self.gen_, h


class Alice:
    def __init__(self, bloodtype, elgamal):
        self.elgamal = elgamal

        self.sk = elgamal.gen_key()

        self.pk = elgamal.gen(self.sk)

        self.b = convert_bloodtype_to_index[convert_from_string_to_enum[bloodtype]]
        self.fake_pks = [self.elgamal.ogen()
                         for _ in range(7)]

    def send_pks(self):
        all_pks = self.fake_pks
        all_pks.insert(self.b-1, self.pk)
        return all_pks

    def retrieve(self, ciphers):
        mb = self.elgamal.dec(ciphers[self.b-1])
        return mb


class Bob:
    def __init__(self, bloodtype, elgamal):
        self.bloodtype = convert_from_string_to_enum[bloodtype]
        self.truth_vals = []
        self.elgamal = elgamal
        self.pks = None
        for donor in BloodType:
            truth_val = blood_cell_compatibility_lookup(self.bloodtype, donor)
            self.truth_vals.append(truth_val)

    def receive_pks(self, pks):
        self.pks = pks

    def transfer_messages(self):
        ciphers = []
        for idx, truth_val in enumerate(self.truth_vals):
            pk = self.pks[idx]
            c = self.elgamal.enc(truth_val, pk)
            ciphers.append(c)
        return ciphers


if __name__ == "__main__":
    p = 199
    q = 2*p + 1
    g = SystemRandom().randint(2, q)

    elgamal = ElGamal(g, q, p)
    alice = Alice("B-", elgamal)
    bob = Bob("B-", elgamal)

    bob.receive_pks(alice.send_pks())
    pls = alice.retrieve(bob.transfer_messages())
    print(pls)