# 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 numpy as np
from crypto.week1 import BloodType, 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(self.order, key) != 1:
            key = SystemRandom().randint(1, self.order)
        return key

    def gen(self, sk):
        h = pow(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, self.order)
        g, h = pk

        s = pow(h, r, self.order)
        p = pow(g, r, self.order)
        c = s * m
        return c, p

    def dec(self, c):
        c1, c2 = c
        h = pow(c2, self.sk, self.order)
        m = c1 / h
        return m

    def ogen(self):
        s = SystemRandom().randint(1, self.order)
        h = pow(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)
        # There is 100% a better way to get the index, also, it's to avoid encryption 0. Coincidentally, it's
        # not an issue when we do it, as O- is in the 0th index and this bloodtype can donate to everyone
        # so the decryption bugging, it not an issue
        self.b = list(convert_bloodtype_to_index.keys()).index(bloodtype)+1
        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):
        print(ciphers)
        mb = self.elgamal.dec(ciphers[self.b-1])
        # Bog sends 1 for false, 2 for true, so we have to subtract 1 here
        return mb-1


class Bob:
    def __init__(self, bloodtype, elgamal):
        self.bloodtype = list(convert_bloodtype_to_index.keys()).index(bloodtype)
        self.truth_vals = []
        self.elgamal = elgamal
        self.pks = None
        # Bob needs his row of the truth table, to create the 8 messages
        for donor in BloodType:
            truth_val = blood_cell_compatibility_lookup(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]
            # Bob can't send 0, as it will encrypt to 0
            c = self.elgamal.enc(truth_val+1, pk)
            ciphers.append(c)
        return ciphers


def run(donor : BloodType, recipient : BloodType):
    p = 15485863
    q = 2 * p + 1
    g = SystemRandom().randint(2, q)

    elgamal = ElGamal(g, q, p)
    alice = Alice(donor, elgamal)
    bob = Bob(recipient, elgamal)

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

    return bool(pls)


if __name__ == "__main__":
    #z = run(BloodType.O_POSITIVE, BloodType.A_NEGATIVE)
    #print(z)
    #exit()

    green = 0
    red = 0
    for i, recipient in enumerate(BloodType):
        for j, donor in enumerate(BloodType):
            z = run(donor, recipient)
            lookup = blood_cell_compatibility_lookup(recipient, donor)
            if lookup == z:
                green += 1
            else:
                print(f"'{BloodType(donor).name} -> {BloodType(recipient).name}' should be {lookup}.")
                red += 1
    print("Green:", green)
    print("Red  :", red)