#include <inttypes.h>
#include <math.h>
#include <pthread.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <time.h>

#include "helper.c"

uint64_t *dec_to_bin(uint64_t d, uint64_t *length) {
    uint64_t *binary_form = calloc(100, sizeof(uint64_t));
    int index = 0;
    while (d != 0) {
        binary_form[index] = d % 2;
        d /= 2;
        index++;
    }

    *length = index;

    return binary_form;
}

uint64_t quick_pow(uint64_t *d_binary, uint64_t a, uint64_t n, uint64_t length) {
    uint64_t *powed = calloc(100, sizeof(uint64_t));

    powed[0] = a;
    for (int i = 1; i <= length; i++) {
        powed[i] = (uint64_t)(((unsigned __int128)powed[i - 1] * powed[i - 1]) % n);
        // printf("powed: %ju, index: %d; ", powed[i], (i));
    }

    // check where in the binary are ones
    uint64_t multiplied = 1;
    for (int i = 0; i < length; i++) {
        if (d_binary[i] == 1) {
            multiplied = (uint64_t)(((unsigned __int128)multiplied * powed[i]) % n);
        }
    }

    // printf("\nbm quick math: %ju; %ju ", multiplied, n);

    free(powed);

    return multiplied;
}

bool prime_test(uint64_t n, int a) {
    // printf("\n\nprime test: %ju\n", n);
    // Miller Rabin prime test
    // choose a base: a, which should be a prime so that (n, a) = 1
    // then do 2 rounds of tests provided the first one did not fail
    // 1: a^d =k 1 mod n
    // 2: a^(d * 2^r) =k n-1 mod n
    // d = n-1 / 2^S (where S means how many time did we divide the number till we reached the first odd number)
    // S: see above
    // r = {0,... S-1}

    // this should not happen but just in case
    if (n <= 1) {
        return false;
    }
    if (n == 2) {
        return true;
    }
    if (n % 2 == 0) {
        return false;
    }

    uint64_t d = n - 1;
    uint64_t S = 0;

    while (d % 2 == 0) {
        d = d / 2;
        S++;
    }

    uint64_t r = S - 1; // this stores the number of elements from 0 to S-1

    // round 1
    // 1: a^d =k 1 mod n
    uint64_t length = 0;
    uint64_t *d_binary = dec_to_bin(d, &length);
    unsigned __int128 first_qp_res = quick_pow(d_binary, a, n, length);

    if (first_qp_res == 1) {
        free(d_binary);
        return true;
    }

    // round 2
    //  2: a^(d * 2^r) =k n-1 mod n
    for (int i = 0; i <= r; i++) {
        if (first_qp_res == n - 1) {
            free(d_binary);
            // printf("true\n");
            return true;
        } else {
            first_qp_res = (uint64_t)(((unsigned __int128)first_qp_res * first_qp_res) % n);
        }
    }

    free(d_binary);
    return false;
}

typedef struct {
    int base;
    uint32_t prime;
} prime_test_t;

void *prime_thread_worker(void *arg) {
    prime_test_t *result_ptr = (prime_test_t *)arg;

    do {
        result_ptr->prime = rand32();
        // printf("\nGenerating a new prime number (%p). Candidate: ", result_ptr);
        // printf("%ju", result_ptr->prime);
        // printf("\n");
    } while (!prime_test(result_ptr->prime, result_ptr->base));

    return NULL;
}

typedef struct {
    uint64_t lnko;
    __int128 x;
    __int128 y;
} euklidian_result_t;

euklidian_result_t euklidian_algorigthm_extended(uint64_t a, uint64_t b) {
    __int128 r = a % b, q = a / b, k = 1, xk = 0, yk = 1, next_r;
    __int128 prev_r = b, prev_q, prev_xk = 0, prev_yk = 1, prev_prev_xk = 1, prev_prev_yk = 0;
    euklidian_result_t res = {0, 0, 0};

    while (r != 0) {
        k++;

        prev_q = q;
        q = prev_r / r;

        next_r = prev_r % r;
        prev_r = r;
        r = next_r;

        xk = prev_xk * prev_q + prev_prev_xk;
        prev_prev_xk = prev_xk;
        prev_xk = xk;

        yk = prev_yk * prev_q + prev_prev_yk;
        prev_prev_yk = prev_yk;
        prev_yk = yk;
    }

    __int128 x = k % 2 == 0 ? prev_xk : -prev_xk;
    __int128 y = k % 2 == 0 ? -prev_yk : prev_yk;

    res.lnko = (uint64_t)prev_r;
    res.x = x;
    res.y = y;

    return res;
}

uint64_t kinai_maradek_tetel(uint64_t *m, uint64_t d, prime_test_t *p, prime_test_t *q) {
    // sum(i: 1,2): Ci * Yi * Mi mod M
    // M: P*Q, Mp: M/P, Mq: M/Q
    uint64_t M = (uint64_t)p->prime * q->prime;
    uint64_t Mp = q->prime;
    uint64_t Mq = p->prime;

    // C1: c^(d mod P-1) mod P
    uint64_t temp_exponent = d % (p->prime - 1);
    uint64_t exponent_bin_length = 0;
    uint64_t *exponent_as_binary = dec_to_bin(temp_exponent, &exponent_bin_length);
    uint64_t c1 = quick_pow(exponent_as_binary, *m, p->prime, exponent_bin_length);
    free(exponent_as_binary);

    // C2: c^(d mod Q-1) mod Q
    temp_exponent = d % (q->prime - 1);
    exponent_as_binary = dec_to_bin(temp_exponent, &exponent_bin_length);
    uint64_t c2 = quick_pow(exponent_as_binary, *m, q->prime, exponent_bin_length);
    free(exponent_as_binary);

    euklidian_result_t y = euklidian_algorigthm_extended(Mp, Mq); // in the struct the x will mean the y1 and y will mean the y2

    // if either of them is less a negative number shift them into postive range with with hte modulo
    uint64_t y1_pos;
    if (y.x < 0) {
        y1_pos = p->prime - (uint64_t)(-y.x % p->prime);
    } else {
        y1_pos = (uint64_t)y.x % p->prime;
    }

    uint64_t y2_pos;
    if (y.y < 0) {
        y2_pos = q->prime - (uint64_t)(-y.y % q->prime);
    } else {
        y2_pos = (uint64_t)y.y % q->prime;
    }

    // Apply the modulo between multiplications to prevent 192-bit overflows!
    uint64_t s1 = (uint64_t)((((unsigned __int128)c1 * y1_pos) % M * Mp) % M);
    uint64_t s2 = (uint64_t)((((unsigned __int128)c2 * y2_pos) % M * Mq) % M);
    return (uint64_t)(((unsigned __int128)s1 + s2) % M);
}

uint64_t rsa_encrypt(uint64_t *m, prime_test_t *p, prime_test_t *q, uint64_t *out_e, uint64_t *out_d) {
    uint64_t n = (uint64_t)p->prime * q->prime;
    printf("n: %ju\n", n);

    uint64_t fi_n = (uint64_t)(p->prime - 1) * (q->prime - 1);
    printf("n: %ju\n", fi_n);

    // 2. kulcsgeneralas
    uint64_t e = 65537;
    do {
        e = rand32() % fi_n;
    } while (e <= 1 || !prime_test(e, p->base)); // the p and q base is used everywhere anyways, i wont pass in another arg

    euklidian_result_t calc_d = euklidian_algorigthm_extended(fi_n, e);

    // if either of them is less a negative number shift them into postive range with with hte modulo
    uint64_t d;
    if (calc_d.y < 0) {
        d = fi_n - (uint64_t)(-calc_d.y);
    } else {
        d = (uint64_t)calc_d.y;
    }

    *out_e = e;
    *out_d = d;

    uint64_t length = 0;
    uint64_t *nyenye = dec_to_bin(e, &length);
    uint64_t c = quick_pow(nyenye, *m, n, length);
    free(nyenye);

    printf("\nc: ");
    print_uint128(c);
    return c;
}

int main() {
    int isSignature = 0;
    printf("Please input 0 for Rsa encryption or 1 for Rsa signature: ");
    scanf("%d", &isSignature);
    printf("\n");

    uint64_t m = 0;
    printf("give input for m: ");
    scanf("%ju", &m);
    printf("\n");

    srand(time(NULL));

    uint64_t base = 2;
    pthread_t thread_p, thread_q;
    prime_test_t p = {base, 0};
    prime_test_t q = {base, 0};

    pthread_create(&thread_p, NULL, prime_thread_worker, &p);
    pthread_create(&thread_q, NULL, prime_thread_worker, &q);

    pthread_join(thread_p, NULL);
    pthread_join(thread_q, NULL);
    printf("\n");

    printf("p: %u\n", p.prime);
    printf("q: %u\n", q.prime);

    uint64_t e = 0;
    uint64_t d = 0;

    if (!isSignature) {
        // rsa encryption
        uint64_t c = rsa_encrypt(&m, &p, &q, &e, &d);
        printf("\nkinai maradek tetel:\n");
        unsigned __int128 S = kinai_maradek_tetel(&c, d, &p, &q);
        printf("S: ");
        print_uint128(S);
        printf("\n");
    } else if (isSignature == 1) {
        // rsa signature

        // generate keys
        uint64_t dummy = 2;
        rsa_encrypt(&dummy, &p, &q, &e, &d);
        printf("\n");

        uint64_t signature = kinai_maradek_tetel(&m, d, &p, &q);
        printf("Alairas (Signature): ");
        print_uint128(signature);

        // key verifacation
        uint64_t e_length = 0;
        uint64_t *e_binary = dec_to_bin(e, &e_length);
        uint64_t n = (uint64_t)p.prime * q.prime;

        uint64_t verified_message = quick_pow(e_binary, signature, n, e_length);
        free(e_binary);

        printf("Verified Message: %ju", verified_message);

        if (verified_message == m) {
            printf("\nSignature correct\n");
        } else {
            printf("\nSignature not correct\n");
        }
    }

    return 0;
}