import gmpy2
from gmpy2 import mpz
from xor import xor_custom
from iloczyn import iloczyn

def odwrotnosc(a):
    """
    Znajduje odwrotność elementu w F_{2^8} za pomocą rozszerzonego algorytmu Euklidesa.

    Args:
        a (mpz): Element w F_{2^8}.

    Returns:
        mpz: Odwrotność elementu w F_{2^8}.
    """
    if not isinstance(a, mpz):
        raise TypeError("Argument musi być typu mpz.")
    
    if a == 0:
        raise ValueError("Element 0 nie ma odwrotności w F_{2^8}.")
    
    # Definicja polinomu niezupełnego (AES)
    irreducible = gmpy2.mpz(0x11B)
    
    # Rozszerzony algorytm Euklidesa
    r0, r1 = irreducible, a
    s0, s1 = gmpy2.mpz(0), gmpy2.mpz(1)
    
    while r1 != 0:
        deg_r0 = r0.bit_length() - 1
        deg_r1 = r1.bit_length() - 1
        if deg_r0 < deg_r1:
            # Swap r0 and r1
            r0, r1 = r1, r0
            s0, s1 = s1, s0
            deg_r0, deg_r1 = deg_r1, deg_r0
        shift = deg_r0 - deg_r1
        # r0 = r0 XOR (r1 << shift)
        r0_shifted = r1 << shift
        r0 = xor_custom(r0, r0_shifted)
        # s0 = s0 XOR (s1 << shift)
        s0_shifted = s1 << shift
        s0 = xor_custom(s0, s0_shifted)
    
    if r0 != 1:
        raise ValueError(f"Element {hex(int(a))} nie ma odwrotności w F_{2^8}.")
    
    return s0 & gmpy2.mpz(0xFF)

if __name__ == "__main__":
    # Przykład użycia
    xy_H = gmpy2.mpz(0x57)  # Przykładowy element
    try:
        wynik = odwrotnosc(xy_H)
        print(f"Odwrotność({hex(int(xy_H))}) = {hex(int(wynik))}")
        # Sprawdzenie: iloczyn powinien być 0x01
        ilocz = iloczyn(xy_H, wynik)
        print(f"Sprawdzenie: iloczyn({hex(int(xy_H))}, {hex(int(wynik))}) = {hex(int(ilocz))}")
    except ValueError as e:
        print(e)