import random
import time
from builtins import int
from typing import Dict, List

size: int = 0


class Country:
    secret: int
    restoredSecret: int
    sharesCount: int
    threshold: int
    secretShares = Dict[int, int]

    def __init__(self, secret: int, sharesCount: int, threshold: int) -> None:
        self.secret = secret
        self.sharesCount = sharesCount
        self.threshold = threshold


def isPrime(n: int) -> bool:    # Test czy wylosowana liczba jest liczba pierwsza
    """
    Miller-Rabin primality test.

    A return value of False means n is certainly not prime. A return value of
    True means n is very likely a prime.
    """
    if n != int(n):
        return False
    n = int(n)
    # Miller-Rabin test for prime
    if n == 0 or n == 1 or n == 4 or n == 6 or n == 8 or n == 9:
        return False

    if n == 2 or n == 3 or n == 5 or n == 7:
        return True
    s = 0
    d = n - 1
    while d % 2 == 0:
        d >>= 1
        s += 1
    assert (2 ** s * d == n - 1)

    def trial_composite(a):
        if pow(a, d, n) == 1:
            return False
        for i in range(s):
            if pow(a, 2 ** i * d, n) == n - 1:
                return False
        return True

    for i in range(8):  # number of trials
        a = random.randrange(2, n)
        if trial_composite(a):
            return False

    return True


def egcd(a, b):
    if a == 0:
        return b, 0, 1
    g, y, x = egcd(b % a, a)
    return g, x - (b // a) * y, y


def modinv(a, m):   # Inversja modulo
    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('No modular inverse')
    return x % m


def getProbablePrime(size: int):
    while True:
        prime = random.getrandbits(size)    # Generuj liczbe o podanym rozmiarze dopki nie bedzie to liczba pierwsza
        if isPrime(prime):                  # o zadanej wielkosci
            return prime


def shareSecret(country: Country, p: int) -> Dict[int, int]:    # Tworzenie cieni
    A: List[int] = list()

    for i in range(0, country.threshold - 1):
        A.append(random.getrandbits(size))

    allShares: Dict[int, int] = dict()

    for i in range(1, country.sharesCount + 1):
        v: int = 0
        for j in range(0, country.threshold - 1):
            v = v + (A[j] * (i ** (country.threshold - j - 1)))
        v = v + country.secret
        allShares[i] = v % p
        print(f"\t{i} : {allShares[i]}")

    return allShares


def pickShares(country: Country) -> Dict[int, int]:     # Wybranie cieni aby odtworzyc sekret potrzebnych jest tyle
    givenSharesIdx: List[int] = list()                  # ile wyznaczonych do odtworzenia jest w danym kraju
    while len(givenSharesIdx) != country.threshold:
        shareNr: int = random.randrange(1, country.sharesCount + 1)
        if shareNr in givenSharesIdx:
            continue
        givenSharesIdx.append(shareNr)

    givenShares: Dict[int, int] = dict()
    for s in country.secretShares.keys():
        if s in givenSharesIdx:
            givenShares[s] = country.secretShares[s]
            # print(f"\t{s}: {givenShares[s]}")
    return givenShares


def restoreSecret(givenShares: Dict[int, int], p: int) -> int:
    secret: int = 0
    for i in givenShares.keys():
        tmp: int = givenShares[i] * delta(i, givenShares, p)    # Odszyfrowanie sekretu
        secret = secret + (tmp % p)                             # Wykorzystujac pod-funckje delta

    return secret % p


def delta(i: int, shares: Dict[int, int], p: int) -> int:   # Tworzenie mini-funkcji dla kazdego z odwiedzanych punktow
    res: int = 1                                            # wspoltworzacych wielomian sekretu
    for s in shares.keys():
        if s != i:
            denominator: int = (i - s) % p
            tmp: int = -s % p
            r: int = tmp % p * modinv(denominator, p)
            r = r % p
            res = res * r
    return res


if __name__ == '__main__':
    random.seed()
    size = int(input("Podaj rozmiar:"))

    p: int = getProbablePrime(size + 5)     # Wyznaczenie wartosci p do operacji modulo
    secret: int = random.getrandbits(size)  # Wyznaczenie wartosci sekretu o zadanej wielkosci

    print("------------- ETAP 1 -------------")
    print("------------- Podział sekretu głównego -------------\n")
    print(f"Sekret to: {secret}")

    countriesCnt: int = 3
    print(f"Dzielę sekret na {countriesCnt} państw")

    secretShares: List[int] = list()
    currShares: int = 0

    for i in range(0, countriesCnt - 1):    # Generowanie sekretow poszczegolnych panstw
        share: int = random.getrandbits(size)
        secretShares.append(share)
        currShares = currShares ^ share
        print(f"Sekret państwa {i + 1} to {share}")

    secretShares.append(currShares ^ secret)

    print(f"Sekret państwa {countriesCnt} to {currShares ^ secret}")

    print("\n------------- ETAP 2 -------------")
    print("------------- Podział sekretu państwa -------------\n")

    countries: List[Country] = list()
    # Lista krajow do odtworzenia sekretu, kraj pierwszy ma 6 senatorow i aby odtworzyc swoj sekret potrzebuje 3
    # senatorow, drugi ma czterech i potrzebuje 3 , trzeci kraj ma 10 i potrzebuje az 6 aby odtworzyc sekret
    countries.append(Country(secretShares[0], 6, 3))
    countries.append(Country(secretShares[1], 4, 3))
    countries.append(Country(secretShares[2], 10, 6))

    for i in range(0, len(countries)):
        x: Country = countries[i]
        print(f"Cienie państwa {i + 1}:")
        x.secretShares = shareSecret(x, p)  # Generowanie cieni dla poszczegolnych panstw

    print("\n------------- ETAP 3 -------------")
    print("------------- Odszyfrowanie sekretów państw -------------\n")

    print("------------- Odszyfrowanie głównego sekretu -------------\n")

    recreationTimes = [0, 0, 0]
    start1 = time.time_ns()
    for _ in range(0, 1000):
        for i in range(0, len(countries)):
            x: Country = countries[i]
            # print(f"Wybór losowych cieni państwa {i + 1}")
            start = time.time_ns()
            givenShares: Dict[int, int] = pickShares(x)     # Wybranie odpowiednich cieni aby odtworzyc sekret
            restoredSecret: int = restoreSecret(givenShares, p)     # Odszyfrowanie sekretu danego Pnastwa
            stop = time.time_ns()
            # print(f"\t\tSekretem państwa {i + 1} jest {restoredSecret}")
            x.restoredSecret = restoredSecret
            recreationTimes[i] += stop - start
    stop1 = time.time_ns()

    print("\n------------- ETAP 4 -------------")
    print("------------- Odszyfrowanie głównego sekretu -------------\n")

    result: int = 0
    for i in range(0, len(countries)):
        result = result ^ countries[i].restoredSecret   # otrzymanie sekretu glownego z sekretow poszczegolnych krajow
        print(f"Odtwarzanie sekretu dla kraju nr {i + 1} : {recreationTimes[i]/(10**6)}")
    print(f"Czas wspolny odtwarzania: {(stop1 - start1)/(10**6)}")
    print(f"Odszyfrowany sekret to {result}")