#include "assert.h"
#include <ch32fun.h>
#include <rand.h>
#include <rsa.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#define LED_PIN PD6
#define RANDOM
#define W 16

void exit_blink() {
    for (int i = 0; i < 4; i++) {
        funDigitalWrite(LED_PIN, FUN_HIGH);
        Delay_Ms(50);
        funDigitalWrite(LED_PIN, FUN_LOW);
        Delay_Ms(50);
    }
}

void enter_blink() {
    for (int i = 0; i < 2; i++) {
        funDigitalWrite(LED_PIN, FUN_HIGH);
        Delay_Ms(200);
        funDigitalWrite(LED_PIN, FUN_LOW);
        Delay_Ms(200);
    }
}

void test_mulmod() {
    ASSERT_EQ(mulmod(3, 2, 4), 2);
    ASSERT_EQ((3 * 2) % 4, 2);

    ASSERT_EQ(mulmod(31, 3, 8), 5);
    ASSERT_EQ(mulmod((u64)1 << 63, 2, 1000000007ULL), 582344008);
}

void test_modexp() {
    ASSERT_EQ(modexp(3, 2, 4), 1);
    ASSERT_EQ((3 ^ 2) % 4, 1);

    ASSERT_EQ(modexp(31, 3, 8), 7);
    ASSERT_EQ(modexp((u64)1 << 63, 2, 1000000007ULL), 319908071);
}

void debug_string(char *str) {
    printf("Got string: %s\n", str);
    for (int i = 0; i < strlen(str); i++) {
        printf("decoded[%d] = '%c' (ASCII: %d)\n", i, str[i],
               str[i]); // Print decoded chars and ASCII values
    }
}

int main() {
    SystemInit();
    sprand(0);

    funGpioInitAll();
    funPinMode(LED_PIN, GPIO_Speed_10MHz | GPIO_CNF_OUT_PP);

    enter_blink();

    test_mulmod();
    test_modexp();

    const u64 p = gen_prime(1 << (W - 1), 1 << W);
    printf("P: %u\n", (u32)p);

    u64 qprev = p;
    while (p == qprev) qprev = gen_prime(1 << (W - 1), 1 << W);

    const u64 q = qprev;
    printf("Q: %u\n", (u32)q);

    ASSERT(gcd(p - 1, PUBEXP) == 1);
    ASSERT(gcd(q - 1, PUBEXP) == 1);

    u64 n = p * q;
    printf("N: %u\n", (u32)n);

    u64 phi_n = (p - 1) * (q - 1);
    printf("Phi_N: %u\n", (u32)phi_n);

    u64 d = mod_inverse(PUBEXP, phi_n);
    printf("D: %u\n", (u32)d);
    if (d == 0 || d == 1) {
        printf("Modular inverse not found...");
    }

    ASSERT_EQ(mulmod(PUBEXP, d, phi_n), 1);

    char msg[] = "Hello";
    u64 coded[sizeof(msg)] = {0};
    char decoded[sizeof(msg)] = {0};

    // Encode the message
    for (int i = 0; i < strlen(msg); i++) {
        coded[i] = modexp((u64)msg[i], PUBEXP, n);
    }

    // Decode the message
    for (int i = 0; i < strlen(msg); i++) {
        u64 dec = modexp(coded[i], d, n);
        decoded[i] = dec & 0xFF;
    }

    {
        printf("Message: %s\n", msg);
        printf("Decoded: %s\n", decoded);

        for (int i = 0; i < strlen(msg); i++) {
            printf("coded[%d] = 0x%016lx\n", i, (unsigned long)coded[i]);
        }

        debug_string(decoded);
    }

    // Exit and hang forever
    exit_blink();
    while (1);
}