# 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.
import math
import random
from secrets import SystemRandom

from .week1 import BloodType, blood_cell_compatibility_lookup

bloodtypes = {b: i for i, b in enumerate(BloodType, start=1)}  # we can't encrypt 0, so we have to index from 1


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 math.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)
        self.b = bloodtypes[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):
        #print(ciphers)
        mb = self.elgamal.dec(ciphers[self.b-1])
        # Bob 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 = bloodtypes[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]
            c = self.elgamal.enc(truth_val + 1, pk)  # + 1 since Bob can't send 0, as it will encrypt to 0
            ciphers.append(c)
        return ciphers


def is_prime(n: int, k: int) -> bool:
    """
    Miller-Rabin Primality test.
    Adapted from pseudo-code at https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test.

    :param n: An odd integer to be tested for primality.
    :param k: The number of rounds of testing to perform.
    :return: True if n is 'probably prime', False otherwise if n is composite.
    """
    # write n as 2r·d + 1 with d odd (by factoring out powers of 2 from n − 1)
    d = n - 1
    r = 0
    while d % 2 == 0:
        d >>= 1
        r += 1

    for i in range(k):  # witnessLoop
        continue_wl = False
        a = random.randint(2, n - 2)
        x = pow(a, d, n)
        if x == 1 or x == n - 1:
            continue
        for j in range(r - 1):
            x = pow(x, 2, n)
            if x == n - 1:
                continue_wl = True
                break
        if continue_wl:
            continue
        return False
    return True


def gen_prime(b: int, k: int = 10) -> int:
    """
    Generate strong probable prime by drawing integers at random until one passes the is_prime test.
    Adapted from pseudo-code at https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test.

    :param b: The number of bits of the result.
    :param k: The number of rounds of testing to perform.
    :return: a strong probable prime.
    """
    while True:
        n = random.randint(2**(b-1), (2**b)-1)
        if n % 2 == 0:
            continue
        if is_prime(n, k):
            return n


def run(donor: BloodType, recipient: BloodType):
    p = gen_prime(128)
    q = 2 * p + 1
    g = SystemRandom().randint(2, q)
    #print("p:", p, "q:", q, "g:", g)

    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)


def main():
    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)