#define MAX_BIT_LENGTH 64

#define MAX_FIRST_PRIMES 1000

#define MAX_RABIN_ITERATIONS 3500

#define PUBLIC_EXPONENT_E 65537

std::bitset<MAX_BIT_LENGTH> Helper::generatePrimeCandidate() {
    // Create bitset of size CANDIDATE_SIZE
    std::bitset<MAX_BIT_LENGTH> candidate_bitset;

    // Set the first and last bit to 1
    candidate_bitset.set(0);
    candidate_bitset.set(MAX_BIT_LENGTH - 1);

    // Generate random number
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, 1);

    // Generate random bits
    for (int i = 1; i < MAX_BIT_LENGTH - 1; i++) { candidate_bitset.set(i, dis(gen)); }

    return candidate_bitset;
}

boost::multiprecision::cpp_int Helper::bitsetToCppInt(std::bitset<MAX_BIT_LENGTH> bitset) {
    return bitset.to_ullong();
}

void Helper::sieveOfEratosthenes(Helper &helper, const size_t max_n) {
    // Create a boolean array "prime[0...n]"
    bool prime[max_n + 1];
    // Initialize all entries it as true. A value in prime[i] will finally be false if i is Not a prime, else true.
    memset(prime, true, sizeof(prime));

    // Loop to update prime[]
    for (int p = 2; p * p <= max_n; p++) {
        // If prime[p] is not changed, then it is a prime
        if (prime[p]) {
            for (int i = p * p; i <= max_n; i += p) { prime[i] = false; }
        }
    }

    // Store primes in FIRST_PRIMES vector
    for (int p = 2; p <= max_n; p++) {
        if (prime[p]) { helper.FIRST_PRIMES.emplace_back(p); }
    }
}

bool RSA::isLowLevelPrimeCandidate(Helper &helper, const boost::multiprecision::cpp_int& candidate) {
    // Get first n primes from helper class
    const std::vector<unsigned int> &FIRST_PRIMES = helper.getFirstPrimes();

    // Loop through first n primes
    for (const auto &FIRST_PRIME: FIRST_PRIMES) {
        // Check if prime candidate is divisible by first n primes
        if (candidate % FIRST_PRIME == 0) return false;
    }

    // Additional checks
    if (candidate <= 1 || candidate % 2 == 0 || candidate % 3 == 0 || candidate % 5 == 0) return false;

    return true;
}

bool RSA::isHighLevelPrimeCandidate(const boost::multiprecision::cpp_int &candidate) {
    boost::multiprecision::cpp_int d = candidate - 1;

    // Find the highest power of 2, k, such that d = 2^k * m
    int k = 0;
    while (d % 2 == 0) {
        d /= 2;
        ++k;
    }
    // Compute m = (n - 1) / 2^k
    boost::multiprecision::cpp_int m = d;

    // Generate random number in range [2, n - 2]
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<unsigned long long> dis(2, static_cast<unsigned long long>(candidate - 2));
    boost::multiprecision::cpp_int a = dis(gen);

    // Compute x = a^m mod n
    boost::multiprecision::cpp_int x = boost::multiprecision::powm(a, m, candidate);

    // Check if x is not 1 and not n - 1
    if (x != 1 && x != candidate - 1) {
        for (boost::multiprecision::cpp_int i = 2; i < MAX_RABIN_ITERATIONS; ++i) {
            x = boost::multiprecision::powm(x, i, candidate);

            if (x == 1) return false; // Composite
            if (x == candidate - 1) return true; // Probable prime
        }

        return false;
    }

    // If x is 1 or n - 1, consider it a probable prime
    return true;
}

boost::multiprecision::cpp_int RSADecryptor::modInverse() {
    return extendedEuclideanAlgorithm(this->getPublicExponentE(), this->getPhiN());
}

boost::multiprecision::cpp_int RSADecryptor::extendedEuclideanAlgorithm(boost::multiprecision::cpp_int a, boost::multiprecision::cpp_int m) {
    // Initialize variables for the extended Euclidean algorithm
    boost::multiprecision::cpp_int original_m = m;
    boost::multiprecision::cpp_int quotient, remainder;
    boost::multiprecision::cpp_int x_current = 0, x_previous = 1;

    // Special case: If m is 1, return 0
    if (m == 1)
        return 0;

    // Apply the extended Euclidean algorithm to find the modular inverse
    while (a > 1) {
        // Calculate quotient and remainder in each iteration
        quotient = a / m;
        remainder = a % m;

        // Update a and m for the next iteration
        a = m;
        m = remainder;

        // Update x_current and x_previous
        boost::multiprecision::cpp_int temp = x_current;
        x_current = x_previous - quotient * x_current;
        x_previous = temp;
    }

    // Make x_previous positive
    if (x_previous < 0)
        x_previous += original_m;

    return x_previous;
}

void Controller::handleUserParameters(int argc, char *argv[]) {
    std::vector<std::string> allowed = {"-h", "--help", "-e", "--encrypt", "-d", "--decrypt"};

    // Check if no arguments are provided
    if (argc == 1) {
        std::cerr << "Error: No arguments provided." << std::endl;
        exit(1);
    }


    // Start at 1 to skip the program name
    for (int i = 1; i < argc; i++) {
        std::string arg = argv[i];

        // Check if argument is allowed
        if (std::find(allowed.begin(), allowed.end(), arg) == allowed.end()) {
            std::cerr << "Error: Invalid argument: " << arg << std::endl;
            exit(1);
        }

        if (arg == "-h" || arg == "--help") {
            Menu::printHelpMenu();
            exit(1);
        }
        else if (arg == "-e" || arg == "--encrypt") {
            if (i + 1 < argc) {
                this->controlEncryption(argv[i + 1]);
                exit(1);
            }
            else {
                std::cerr << "Error: -e option requires one argument." << std::endl;
                exit(1);
            }
        }
        else if (arg == "-d" || arg == "--decrypt") {
            if (i + 3 < argc) {
                // Convert private key and modulus N to cpp_int
                try {
                    boost::multiprecision::cpp_int privateKeyD = helper.stringToCppInt(argv[i + 2]);
                    boost::multiprecision::cpp_int modulusN = helper.stringToCppInt(argv[i + 3]);
                    // Decrypt with additional parameters
                    this->controlDecryption(argv[i + 1], privateKeyD, modulusN);
                    i += 3; // Skip the additional parameters
                    exit(1);
                } catch (const std::exception& e) {
                    std::cerr << "\nError: -d option requires three arguments: <string> <private_key> <modulus_N>"
                              << std::endl;
                    exit(1);
                }
            } else {
                std::cerr << "\nError: -d option requires three arguments: <string> <private_key> <modulus_N>"
                          << std::endl;
                exit(1);
            }
        }
    }
}

void Controller::controlEncryption(const std::string& message) {
    // Generate prime numbers
    helper.sieveOfEratosthenes(helper);

    // ======== Generate prime candidates ========
    PrimeCandidate p_candidate = PrimeCandidate();
    PrimeCandidate q_candidate = PrimeCandidate();

    while (true) {
        // Generate prime candidate
        std::bitset<MAX_BIT_LENGTH> p_candidate_bits = helper.generatePrimeCandidate();

        // Convert prime candidate to cpp_int
        p_candidate.setCandidate(helper.bitsetToCppInt(p_candidate_bits));

        // Check if prime candidate is prime
        if (rsa.isLowLevelPrimeCandidate(helper, p_candidate.getCandidate()) && rsa.isHighLevelPrimeCandidate(p_candidate.getCandidate())) {
            std::cout << "Prime candidate P is: " << p_candidate.getCandidate() << std::endl;
            p_candidate.setIsPrime(true);
            break;
        }
    }

    while (true) {
        // Generate prime candidate
        std::bitset<MAX_BIT_LENGTH> q_candidate_bits = helper.generatePrimeCandidate();

        // Convert prime candidate to cpp_int
        q_candidate.setCandidate(helper.bitsetToCppInt(q_candidate_bits));

        // Check if prime candidate is prime
        if (rsa.isLowLevelPrimeCandidate(helper, q_candidate.getCandidate()) && rsa.isHighLevelPrimeCandidate(q_candidate.getCandidate())) {
            std::cout << "Prime candidate Q is: " << q_candidate.getCandidate() << std::endl;
            q_candidate.setIsPrime(true);
            break;
        }
    }


    // Additional check to ensure that P and Q are not equal
    if (p_candidate.getCandidate() == q_candidate.getCandidate()) {
        std::cout << "P and Q are equal, generating new prime candidates..." << std::endl;
        controlEncryption(message);
    }

    // ======== Compute modulus N, phi N and public exponent E ========
    if (p_candidate.getIsPrime() && q_candidate.getIsPrime()) {
        rsa.computeModulusN(p_candidate.getCandidate(), q_candidate.getCandidate());
        rsa.computePhiN(p_candidate.getCandidate(), q_candidate.getCandidate());
        rsa.setPublicExponentE();
    }

    std::cout << "Modulus N is: " << rsa.getModulusN() << std::endl;
    std::cout << "Phi N is: " << rsa.getPhiN() << std::endl;
    std::cout << "Public exponent E is: " << rsa.getPublicExponentE() << std::endl;

    // Encrypt message
    std::cout << "Message to encrypt is: " << message << std::endl;

    // ======== Encryption ========
    RSAEncryptor rsaEncryptor = RSAEncryptor(rsa, message);

    auto e_start = std::chrono::high_resolution_clock::now();
    std::vector<boost::multiprecision::cpp_int> encrypted =  rsaEncryptor.encryptMessage();
    auto e_end = std::chrono::high_resolution_clock::now();

    std::cout << "Encrypted message is: ";
    for (auto & i : encrypted) std::cout << i << " ";
    std::cout << std::endl;

    std::cout << "Encryption took: " << to_ms(e_end - e_start).count() << "ms" << std::endl;

    std::cout << std::endl;

    // ======== Decryption ========
    std::cout << "Initializing testing decryption..." << std::endl;

    // Decrypt message - compute private exponent D
    RSADecryptor rsaDecryptor = RSADecryptor(rsa);

    // Compute private exponent D
    rsaDecryptor.findPrivateKeyD();

    std::cout << "Private exponent D is: " << rsaDecryptor.getPrivateKeyD() << std::endl;

    auto d_start = std::chrono::high_resolution_clock::now();
    // Decrypt message
    std::string decrypted = rsaDecryptor.decryptMessage(encrypted);
    auto d_end = std::chrono::high_resolution_clock::now();

    std::cout << "Decrypted message is: " << decrypted << std::endl;
    std::cout << std::endl;

    std::cout << "Decryption took: " << to_ms(d_end - d_start).count() << "ms" << std::endl;

    // ======== Outro ========
    std::cout << std::endl;
    std::cout << "For decryption of your encrypted message, use the following parameters: " << std::endl;
    std::cout << "Encrypted message: ";
    for (auto & i : encrypted) std::cout << i << " ";
    std::cout << std::endl;
    std::cout << "Private exponent D: " << rsaDecryptor.getPrivateKeyD() << std::endl;
    std::cout << "Modulus N: " << rsa.getModulusN() << std::endl;
    std::cout << std::endl;
}

void Controller::controlDecryption(const std::string &message, const boost::multiprecision::cpp_int &privateKeyD,
                                   const boost::multiprecision::cpp_int &modulusN) {
    std::cout << "Initializing decryption..." << std::endl;

    auto d_start = std::chrono::high_resolution_clock::now();
    std::string decrypted = RSADecryptor::decryptMessage(message, privateKeyD, modulusN);
    auto d_end = std::chrono::high_resolution_clock::now();

    std::cout << "Decrypted message is: " << decrypted << std::endl;

    std::cout << "Decryption took: " << to_ms(d_end - d_start).count() << "ms" << std::endl;

    std::cout << std::endl;
}