from __future__ import annotations


import random

from .week1 import BloodType, blood_cell_compatibility_lookup


class Protocol:
    def __init__(self, alice, bob):
        self.alice: Alice = alice
        self.bob: Bob = bob
        self.dealer = Dealer()

        self.alice.share_inputs(bob)
        self.bob.share_inputs(alice)

    def xor(self, x_name, y_name, z_name):
        self.alice.xor(x_name, y_name, z_name)
        self.bob.xor(x_name, y_name, z_name)

    def xor_with_constant(self, x_name, c, z_name):
        self.alice.xor_with_constant(x_name, c, z_name)
        self.bob.xor_with_constant(x_name, c, z_name)

    def and_(self, x_name, y_name, z_name):
        self.dealer.gen_randoms()
        self.alice.set_randoms(self.dealer.randoms_alice)
        self.bob.set_randoms(self.dealer.randoms_bob)

        self.xor(x_name, "u", "d")
        self.xor(y_name, "v", "e")

        self.open("d")
        self.open("e")

        # z = w ^ (e & x) ^ (d & y) ^ (e & d)
        # z = w ^    a    ^    b    ^    c
        # z =   z1        ^         z2
        self.and_with_public(x_name, "e", "a")
        self.and_with_public(y_name, "d", "b")
        self.and_publics("e", "d", "c")

        self.xor("w", "a", "z1")
        self.xor_with_public("b", "c", "z2")

        self.xor("z1", "z2", z_name)

    def and_with_public(self, x_name, y_name, z_name):
        self.alice.and_with_public(x_name, y_name, z_name)
        self.bob.and_with_public(x_name, y_name, z_name)

    def and_publics(self, x_name, y_name, z_name):
        self.alice.and_publics(x_name, y_name, z_name)
        self.bob.and_publics(x_name, y_name, z_name)

    def xor_with_public(self, x_name, c_name, z_name):
        self.alice.xor_with_public(x_name, c_name, z_name)
        self.bob.xor_with_public(x_name, c_name, z_name)

    def open(self, x_name):
        self.alice.open(x_name, self.bob.send(x_name))
        self.bob.open(x_name, self.alice.send(x_name))

    def output(self, z_name):
        return self.alice.open(z_name, self.bob.send(z_name))


class Party:
    def __init__(self, **inputs) -> None:
        self.inputs = inputs
        self.secrets = {}
        self.publics = {}

    def share(self, x_name):
        x_other = random.randint(0, 1)
        self.secrets[x_name] = self.inputs[x_name] ^ x_other
        return x_other

    def share_inputs(self, other_party: Party):
        for name, value in self.inputs.items():
            other_party.receive(name, self.share(name))

    def set_randoms(self, randoms):
        self.secrets["u"], self.secrets["v"], self.secrets["w"] = randoms

    def xor(self, x_name, y_name, z_name):
        x = self.secrets[x_name]
        y = self.secrets[y_name]
        self.secrets[z_name] = x ^ y

    def xor_with_constant(self, x_name, c, z_name):
        x = self.secrets[x_name]
        self.secrets[z_name] = c ^ x

    def and_with_public(self, x_name, y_name, z_name):
        x = self.secrets[x_name]
        y = self.publics[y_name]
        self.secrets[z_name] = x & y

    def and_publics(self, x_name, y_name, z_name):
        x = self.publics[x_name]
        y = self.publics[y_name]
        self.publics[z_name] = x & y

    def xor_with_public(self, x_name, c_name, z_name):
        x = self.secrets[x_name]
        c = self.publics[c_name]
        self.secrets[z_name] = c ^ x

    def open(self, x_name, other_share):
        x = self.secrets[x_name]
        z = x ^ other_share
        self.publics[x_name] = z
        return z

    def send(self, x_name):
        return self.secrets[x_name]

    def receive(self, x_name, x):
        self.secrets[x_name] = x


class Alice(Party):
    pass


class Bob(Party):
    def xor_with_constant(self, x_name, c, z_name):
        """
        xor with constant is asymmetric, so bob simply saves x to z.
        """
        x = self.secrets[x_name]
        self.secrets[z_name] = x

    def xor_with_public(self, x_name, c_name, z_name):
        """
        xor with constant is asymmetric, so bob simply saves x to z.
        """
        x = self.secrets[x_name]
        self.secrets[z_name] = x


class Dealer:
    def gen_randoms(self):
        u = random.randint(0,1)
        v = random.randint(0,1)
        w = u & v

        u_b = random.randint(0,1)
        u_a = u_b ^ u

        v_b = random.randint(0,1)
        v_a = v_b ^ v

        w_b = random.randint(0,1)
        w_a = w_b ^ w

        assert((u_a ^ u_b) & (v_a ^ v_b) == w)
        assert(w_a ^ w_b == w)

        self.randoms_alice = (u_a, v_a, w_a)
        self.randoms_bob = (u_b, v_b, w_b)


def run(da, db, ds, ra, rb, rs):
    proto = Protocol(Alice(da=da, db=db, ds=ds), Bob(ra=ra, rb=rb, rs=rs))

    proto.xor_with_constant("ra", 1, "k1")
    proto.xor_with_constant("rb", 1, "k2")
    proto.xor_with_constant("rs", 1, "k3")

    proto.and_("da", "k1", "l1")
    proto.and_("db", "k2", "l2")
    proto.and_("ds", "k3", "l3")

    proto.xor_with_constant("l1", 1, "m1")
    proto.xor_with_constant("l2", 1, "m2")
    proto.xor_with_constant("l3", 1, "m3")

    proto.and_("m1", "m2", "n")
    proto.and_("n", "m3", "z")

    z = proto.output("z")
    return z


def main():
    green = 0
    red = 0
    for i, recipient in enumerate(BloodType):
        for j, donor in enumerate(BloodType):
            z = run(*donor.value, *recipient.value)
            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)