#include "bech32.h"

#include <stddef.h>
#include <stdint.h>
#include <string.h>

static const char BECH32_CHARSET[] = "qpzry9x8gf2tvdw0s3jn54khce6mua7l";

static uint32_t bech32_polymod_step(uint32_t pre)
{
    const uint8_t b = (uint8_t)(pre >> 25);
    return ((pre & 0x1FFFFFFu) << 5) ^
           ((-(int32_t)((b >> 0) & 1u)) & 0x3b6a57b2u) ^
           ((-(int32_t)((b >> 1) & 1u)) & 0x26508e6du) ^
           ((-(int32_t)((b >> 2) & 1u)) & 0x1ea119fau) ^
           ((-(int32_t)((b >> 3) & 1u)) & 0x3d4233ddu) ^
           ((-(int32_t)((b >> 4) & 1u)) & 0x2a1462b3u);
}

static int convert_bits(uint8_t *out,
                        size_t *out_len,
                        int out_bits,
                        const uint8_t *in,
                        size_t in_len,
                        int in_bits,
                        int pad)
{
    uint32_t value = 0;
    int bits = 0;
    const uint32_t max_v = (((uint32_t)1) << out_bits) - 1;

    *out_len = 0;

    while (in_len-- > 0) {
        value = (value << in_bits) | *(in++);
        bits += in_bits;

        while (bits >= out_bits) {
            bits -= out_bits;
            out[(*out_len)++] = (uint8_t)((value >> bits) & max_v);
        }
    }

    if (pad) {
        if (bits > 0) {
            out[(*out_len)++] = (uint8_t)((value << (out_bits - bits)) & max_v);
        }
    } else if (((value << (out_bits - bits)) & max_v) != 0u || bits >= in_bits) {
        return 0;
    }

    return 1;
}

static int bech32_encode(char *output, size_t output_len, const char *hrp, const uint8_t *data, size_t data_len)
{
    uint32_t chk = 1;
    size_t hrp_len = strlen(hrp);
    size_t pos = 0;

    if (output == NULL || hrp == NULL || data == NULL) {
        return 0;
    }

    for (size_t i = 0; i < hrp_len; ++i) {
        const unsigned char ch = (unsigned char)hrp[i];
        if (ch < 33 || ch > 126 || (ch >= 'A' && ch <= 'Z')) {
            return 0;
        }
        chk = bech32_polymod_step(chk) ^ (ch >> 5);
    }

    chk = bech32_polymod_step(chk);

    for (size_t i = 0; i < hrp_len; ++i) {
        chk = bech32_polymod_step(chk) ^ (hrp[i] & 0x1f);
        if (pos + 1 >= output_len) {
            return 0;
        }
        output[pos++] = hrp[i];
    }

    if (pos + 1 >= output_len) {
        return 0;
    }
    output[pos++] = '1';

    for (size_t i = 0; i < data_len; ++i) {
        if ((data[i] >> 5) != 0u) {
            return 0;
        }
        chk = bech32_polymod_step(chk) ^ data[i];
        if (pos + 1 >= output_len) {
            return 0;
        }
        output[pos++] = BECH32_CHARSET[data[i]];
    }

    for (size_t i = 0; i < 6; ++i) {
        chk = bech32_polymod_step(chk);
    }

    chk ^= 1;

    for (size_t i = 0; i < 6; ++i) {
        if (pos + 1 >= output_len) {
            return 0;
        }
        output[pos++] = BECH32_CHARSET[(chk >> ((5 - i) * 5)) & 0x1f];
    }

    if (pos >= output_len) {
        return 0;
    }

    output[pos] = '\0';
    return 1;
}

int npub_encode(const uint8_t pubkey32[32], char *out, size_t out_len)
{
    uint8_t conv[64] = {0};
    size_t conv_len = 0;

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

    if (!convert_bits(conv, &conv_len, 5, pubkey32, 32, 8, 1)) {
        return -1;
    }

    if (!bech32_encode(out, out_len, "npub", conv, conv_len)) {
        return -1;
    }

    memset(conv, 0, sizeof(conv));
    return 0;
}
