from __future__ import annotations

import hashlib
import itertools
import random
import secrets
from typing import List

from week1 import BloodType


def sha256(b: bytes) -> bytes:
    return hashlib.sha256(b).digest()


def rand_bytes():
    return secrets.SystemRandom().getrandbits(128).to_bytes(16, "big")


def xor_bytes(a: bytes, b: bytes, k=256) -> bytes:
    return (int.from_bytes(a, "big") ^ int.from_bytes(b, "big")).to_bytes(k, "big")


class Gate:
    def __init__(self, left: Gate, right: Gate, index: int) -> None:
        self.left = left
        self.right = right
        self.i = index
        self.k = {
            0: rand_bytes(),
            1: rand_bytes()
        }
        self.output = None

        c_prime = {}
        for a, b in itertools.product((0, 1), repeat=2):
            c_prime[(a, b)] = xor_bytes(
                sha256(self.left.k[a] + self.right.k[b] + self.i.to_bytes(1, "big")),
                self.k[not a * b] + bytes(128 // 8)
            )
        pi = list(itertools.product((0, 1), repeat=2))
        random.shuffle(pi)
        self.c = {i: c_prime[p]
                  for i, p in enumerate(pi)}


class InputWire:
    def __init__(self, index) -> None:
        self.i = index
        self.k = {
            0: rand_bytes(),
            1: rand_bytes()
        }
        self.output = None


class Circuit:
    def __init__(self, input_wires: List[InputWire], gates: List[Gate]) -> None:
        self.input_wires = input_wires
        self.gates = gates

    def evaluate(self, x: List[bytes]) -> bytes:
        for i, xi in enumerate(x):
            self.input_wires[i].output = xi

        for gate in self.gates:
            for j in range(4):
                xor = xor_bytes(
                    sha256(gate.left.output + gate.right.output + gate.i),
                    gate.c[j]
                )
                k_prime, tau = xor[:16], xor[16:]
                if tau == bytes(16):
                    gate.output = k_prime

        return self.gates[-1].output


def encode(e: List[InputWire], x: List[int]) -> List[bytes]:
    return [e[i].k[xi]
            for i, xi in enumerate(x)]


class Alice:
    pass


class Bob:
    pass


def run(donor: BloodType, recipient: BloodType):
    m1 = Alice.choose(x)
    m2 = Bob.transfer(y, m2)
    z = Alice.retrieve(m2)
    return z


def main():
    run(donor=BloodType.A_NEGATIVE, recipient=BloodType.B_POSITIVE)