# 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, 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(self.order, 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, self.order)
        g, h = pk

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

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

    def ogen(self):
        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 = 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):
        mb = self.elgamal.dec(ciphers[self.b-1])
        return mb


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
        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]
            c = self.elgamal.enc(truth_val, pk)
            ciphers.append(c)
        return ciphers


def run(donor : BloodType, recipient : BloodType):
    p = 199
    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__":
    # 199 is a large prime

    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)