#include "mnemonic.h"

#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <ctype.h>

#include "esp_random.h"

#include "mbedtls/md.h"
#include "mbedtls/pkcs5.h"
#include "mbedtls/sha256.h"

#include "mnemonic_wordlist.h"

#define MNEMONIC_12_WORDS 12
#define ENTROPY_12_BYTES 16
#define CHECKSUM_12_BITS 4
#define TOTAL_12_BITS (ENTROPY_12_BYTES * 8 + CHECKSUM_12_BITS)

static bool s_letter_index_built = false;
static size_t s_letter_start[26] = {0};
static size_t s_letter_count[26] = {0};

static bool append_word(char *out, size_t out_len, size_t *cursor, const char *word)
{
    const size_t word_len = strlen(word);

    if (*cursor >= out_len) {
        return false;
    }

    if (*cursor != 0) {
        if (*cursor + 1 >= out_len) {
            return false;
        }
        out[*cursor] = ' ';
        (*cursor)++;
    }

    if (*cursor + word_len >= out_len) {
        return false;
    }

    memcpy(out + *cursor, word, word_len);
    *cursor += word_len;
    out[*cursor] = '\0';
    return true;
}

static uint8_t get_bit_from_entropy_checksum(const uint8_t entropy[ENTROPY_12_BYTES],
                                             const uint8_t checksum[32],
                                             size_t bit_index)
{
    if (bit_index < ENTROPY_12_BYTES * 8) {
        const size_t byte_index = bit_index / 8;
        const size_t bit_in_byte = 7 - (bit_index % 8);
        return (uint8_t)((entropy[byte_index] >> bit_in_byte) & 0x01);
    }

    const size_t checksum_bit_index = bit_index - (ENTROPY_12_BYTES * 8);
    const size_t bit_in_byte = 7 - checksum_bit_index;
    return (uint8_t)((checksum[0] >> bit_in_byte) & 0x01);
}

static int bip39_word_index(const char *word)
{
    if (word == NULL || word[0] == '\0') {
        return -1;
    }

    for (size_t i = 0; i < 2048; ++i) {
        if (strcmp(word, BIP39_WORDLIST[i]) == 0) {
            return (int)i;
        }
    }

    return -1;
}

static bool next_token(const char **cursor, char *out_word, size_t out_len)
{
    size_t n = 0;
    const char *p = *cursor;

    while (*p != '\0' && isspace((unsigned char)*p)) {
        ++p;
    }

    if (*p == '\0') {
        *cursor = p;
        return false;
    }

    while (*p != '\0' && !isspace((unsigned char)*p)) {
        if (n + 1 >= out_len) {
            return false;
        }
        out_word[n++] = *p++;
    }

    out_word[n] = '\0';
    *cursor = p;
    return true;
}

static void build_letter_index(void)
{
    if (s_letter_index_built) {
        return;
    }

    for (size_t i = 0; i < 26; ++i) {
        s_letter_start[i] = 0;
        s_letter_count[i] = 0;
    }

    for (size_t i = 0; i < 2048; ++i) {
        char c = BIP39_WORDLIST[i][0];
        if (c < 'a' || c > 'z') {
            continue;
        }

        size_t li = (size_t)(c - 'a');
        if (s_letter_count[li] == 0) {
            s_letter_start[li] = i;
        }
        s_letter_count[li]++;
    }

    s_letter_index_built = true;
}

int generate_mnemonic_12(char *out, size_t out_len)
{
    uint8_t entropy[ENTROPY_12_BYTES] = {0};
    uint8_t hash[32] = {0};
    size_t cursor = 0;

    if (out == NULL || out_len == 0) {
        return -1;
    }

    out[0] = '\0';

    esp_fill_random(entropy, sizeof(entropy));
    mbedtls_sha256(entropy, sizeof(entropy), hash, 0);

    for (size_t w = 0; w < MNEMONIC_12_WORDS; ++w) {
        const size_t start_bit = w * 11;
        uint16_t index = 0;

        for (size_t b = 0; b < 11; ++b) {
            index = (uint16_t)((index << 1) |
                               get_bit_from_entropy_checksum(entropy, hash, start_bit + b));
        }

        if (index >= 2048) {
            memset(entropy, 0, sizeof(entropy));
            memset(hash, 0, sizeof(hash));
            return -1;
        }

        if (!append_word(out, out_len, &cursor, BIP39_WORDLIST[index])) {
            memset(entropy, 0, sizeof(entropy));
            memset(hash, 0, sizeof(hash));
            out[0] = '\0';
            return -1;
        }
    }

    if (TOTAL_12_BITS != 132) {
        memset(entropy, 0, sizeof(entropy));
        memset(hash, 0, sizeof(hash));
        out[0] = '\0';
        return -1;
    }

    memset(entropy, 0, sizeof(entropy));
    memset(hash, 0, sizeof(hash));
    return 0;
}

int mnemonic_validate(const char *mnemonic)
{
    uint16_t word_indices[MNEMONIC_12_WORDS] = {0};
    uint8_t entropy[ENTROPY_12_BYTES] = {0};
    uint8_t hash[32] = {0};
    const char *cursor = mnemonic;
    char token[16];
    size_t words = 0;
    size_t bit_pos = 0;

    if (mnemonic == NULL) {
        return -1;
    }

    while (next_token(&cursor, token, sizeof(token))) {
        int idx;

        if (words >= MNEMONIC_12_WORDS) {
            memset(entropy, 0, sizeof(entropy));
            memset(hash, 0, sizeof(hash));
            return -1;
        }

        idx = bip39_word_index(token);
        if (idx < 0) {
            memset(entropy, 0, sizeof(entropy));
            memset(hash, 0, sizeof(hash));
            return -1;
        }

        word_indices[words++] = (uint16_t)idx;
    }

    if (words != MNEMONIC_12_WORDS) {
        memset(entropy, 0, sizeof(entropy));
        memset(hash, 0, sizeof(hash));
        return -1;
    }

    for (size_t w = 0; w < MNEMONIC_12_WORDS; ++w) {
        for (int b = 10; b >= 0; --b) {
            uint8_t bit = (uint8_t)((word_indices[w] >> b) & 0x01);
            if (bit_pos < ENTROPY_12_BYTES * 8) {
                size_t byte_index = bit_pos / 8;
                size_t bit_in_byte = 7 - (bit_pos % 8);
                entropy[byte_index] |= (uint8_t)(bit << bit_in_byte);
            }
            bit_pos++;
        }
    }

    mbedtls_sha256(entropy, sizeof(entropy), hash, 0);

    for (size_t i = 0; i < CHECKSUM_12_BITS; ++i) {
        size_t checksum_bit_pos = ENTROPY_12_BYTES * 8 + i;
        size_t word_idx = checksum_bit_pos / 11;
        size_t bit_in_word = 10 - (checksum_bit_pos % 11);
        uint8_t expected = (uint8_t)((hash[0] >> (7 - i)) & 0x01);
        uint8_t actual = (uint8_t)((word_indices[word_idx] >> bit_in_word) & 0x01);
        if (actual != expected) {
            memset(entropy, 0, sizeof(entropy));
            memset(hash, 0, sizeof(hash));
            return -1;
        }
    }

    memset(entropy, 0, sizeof(entropy));
    memset(hash, 0, sizeof(hash));
    return 0;
}

size_t mnemonic_words_for_letter(char letter, size_t *out_start)
{
    int idx;

    if (out_start == NULL) {
        return 0;
    }

    if (letter >= 'A' && letter <= 'Z') {
        letter = (char)(letter - 'A' + 'a');
    }

    if (letter < 'a' || letter > 'z') {
        *out_start = 0;
        return 0;
    }

    build_letter_index();

    idx = letter - 'a';
    *out_start = s_letter_start[idx];
    return s_letter_count[idx];
}

const char *mnemonic_word_at(size_t index)
{
    if (index >= 2048) {
        return NULL;
    }
    return BIP39_WORDLIST[index];
}

int mnemonic_to_seed(const char *mnemonic, uint8_t seed[64])
{
    const mbedtls_md_info_t *md_info;
    mbedtls_md_context_t ctx;
    uint8_t salt[8] = {'m', 'n', 'e', 'm', 'o', 'n', 'i', 'c'};

    if (mnemonic == NULL || seed == NULL) {
        return -1;
    }

    md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA512);
    if (md_info == NULL) {
        return -1;
    }

    mbedtls_md_init(&ctx);
    if (mbedtls_md_setup(&ctx, md_info, 1) != 0) {
        mbedtls_md_free(&ctx);
        return -1;
    }

    (void)ctx;
    if (mbedtls_pkcs5_pbkdf2_hmac_ext(MBEDTLS_MD_SHA512,
                                      (const uint8_t *)mnemonic,
                                      strlen(mnemonic),
                                      salt,
                                      sizeof(salt),
                                      2048,
                                      64,
                                      seed) != 0) {
        mbedtls_md_free(&ctx);
        return -1;
    }

    mbedtls_md_free(&ctx);
    return 0;
}