#include "PrecomputedPuzzles.h"
#include "../../Configuration.h"
#include "../ConnectedClient.h"
#include <tomcrypt.h>

using namespace std;
using namespace ts;
using namespace ts::protocol;

PuzzleManager::PuzzleManager() {}
PuzzleManager::~PuzzleManager() {}

size_t PuzzleManager::precomputedPuzzleCount() { return this->cached.size(); }

bool PuzzleManager::precomputePuzzles(size_t limit) {
    while(precomputedPuzzleCount() < limit) generatePuzzle();
    return true;
}

std::shared_ptr<Puzzle> PuzzleManager::nextPuzzle() {
    this->indexLock.lock();
    size_t index = this->cacheIndex++ % this->cached.size();
    this->indexLock.unlock();
    return this->cached[index];
}

inline void rndNum(mp_int *result, int byteLength){
    uint8_t buffer[byteLength];

    for(int index = 0; index < byteLength; index++) {
        int rnd = rand();
        uint8_t urnd = static_cast<uint8_t>(rnd & 0xFF);
        buffer[index] = urnd; //TODO more secure!

    }

    mp_zero(result);
    mp_read_unsigned_bin(result, buffer, byteLength);
}

inline bool solvePuzzle(Puzzle *puzzle){
    mp_int exp{};
    mp_init(&exp);
    mp_2expt(&exp, puzzle->level);


    if (mp_exptmod(&puzzle->x, &exp, &puzzle->n, &puzzle->result) != CRYPT_OK) { //Sometimes it fails (unknow why :D)
        mp_clear(&exp);
        return false;
    }

    mp_clear(&exp);
    return true;
}

inline bool write_bin_data(mp_int& data, uint8_t* result, size_t length) {
    ulong n{length};
    if(auto err = mp_to_unsigned_bin_n(&data, result, &n); err)
        return false;

    if(n != length) {
        auto off = length - n;
        memcpy(result + off, result, n);
        memset(result, 0, off);
    }
    return true;
}

void PuzzleManager::generatePuzzle() {
    auto puzzle = new Puzzle{};

    mp_init_multi(&puzzle->x, &puzzle->n, &puzzle->result, nullptr);

    generate_new:
    rndNum(&puzzle->x, 64);
    rndNum(&puzzle->n, 64);
    puzzle->level = ts::config::voice::RsaPuzzleLevel;

    if(!solvePuzzle(puzzle))
        goto generate_new;

    auto valid_x = mp_unsigned_bin_size(&puzzle->x) <= 64;
    auto valid_n = mp_unsigned_bin_size(&puzzle->n) <= 64;
    auto valid_result = mp_unsigned_bin_size(&puzzle->result) <= 64;
    if(!valid_n || !valid_x || !valid_result)
        goto generate_new;

    if(!write_bin_data(puzzle->x, puzzle->data_x, 64))
        goto generate_new;

    if(!write_bin_data(puzzle->n, puzzle->data_n, 64))
        goto generate_new;

    if(!write_bin_data(puzzle->result, puzzle->data_result, 64))
        goto generate_new;

    this->cached.push_back(shared_ptr<Puzzle>(puzzle, [](Puzzle* elm){
        mp_clear_multi(&elm->n, &elm->x, &elm->result, nullptr);
        delete elm;
    }));
}