// SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
// SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later
// Licensed under OWL v1.0+. See LICENSE.
/*
 * LITLS -- Base64 and Base64url encode/decode (RFC 4648)
 */

#include "litls.h"

static const char b64_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

static const char b64url_table[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

static size_t encode_with_table(const uint8_t *in, size_t in_len, char *out, size_t out_cap, const char *table,
                                int pad) {
    size_t needed = ((in_len + 2) / 3) * 4;
    if (!pad) {
        size_t rem = in_len % 3;
        if (rem == 1)
            needed -= 2;
        else if (rem == 2)
            needed -= 1;
    }
    if (out_cap < needed + 1)
        return 0;

    size_t i = 0, j = 0;
    while (i + 2 < in_len) {
        uint32_t v = ((uint32_t)in[i] << 16) | ((uint32_t)in[i + 1] << 8) | in[i + 2];
        out[j++]   = table[(v >> 18) & 0x3f];
        out[j++]   = table[(v >> 12) & 0x3f];
        out[j++]   = table[(v >> 6) & 0x3f];
        out[j++]   = table[v & 0x3f];
        i += 3;
    }

    if (i < in_len) {
        uint32_t v = (uint32_t)in[i] << 16;
        if (i + 1 < in_len)
            v |= (uint32_t)in[i + 1] << 8;
        out[j++] = table[(v >> 18) & 0x3f];
        out[j++] = table[(v >> 12) & 0x3f];
        if (i + 1 < in_len) {
            out[j++] = table[(v >> 6) & 0x3f];
            if (pad)
                out[j++] = '=';
        } else {
            if (pad) {
                out[j++] = '=';
                out[j++] = '=';
            }
        }
    }

    out[j] = '\0';
    return j;
}

static int decode_char(char c, const char *table) {
    for (int i = 0; i < 64; i++) {
        if (table[i] == c)
            return i;
    }
    return -1;
}

static int decode_with_table(const char *in, size_t in_len, uint8_t *out, size_t *out_len, const char *table) {
    /* strip trailing padding */
    while (in_len > 0 && in[in_len - 1] == '=')
        in_len--;

    size_t max_out = (in_len * 3) / 4;
    if (*out_len < max_out)
        return -1;

    size_t i = 0, j = 0;
    while (i + 3 < in_len) {
        int a = decode_char(in[i], table);
        int b = decode_char(in[i + 1], table);
        int c = decode_char(in[i + 2], table);
        int d = decode_char(in[i + 3], table);
        if (a < 0 || b < 0 || c < 0 || d < 0)
            return -1;
        out[j++] = (uint8_t)((a << 2) | (b >> 4));
        out[j++] = (uint8_t)(((b & 0x0f) << 4) | (c >> 2));
        out[j++] = (uint8_t)(((c & 0x03) << 6) | d);
        i += 4;
    }

    if (in_len - i == 3) {
        int a = decode_char(in[i], table);
        int b = decode_char(in[i + 1], table);
        int c = decode_char(in[i + 2], table);
        if (a < 0 || b < 0 || c < 0)
            return -1;
        out[j++] = (uint8_t)((a << 2) | (b >> 4));
        out[j++] = (uint8_t)(((b & 0x0f) << 4) | (c >> 2));
    } else if (in_len - i == 2) {
        int a = decode_char(in[i], table);
        int b = decode_char(in[i + 1], table);
        if (a < 0 || b < 0)
            return -1;
        out[j++] = (uint8_t)((a << 2) | (b >> 4));
    } else if (in_len - i != 0) {
        return -1;
    }

    *out_len = j;
    return 0;
}

size_t litls_base64_encode(const uint8_t *in, size_t in_len, char *out, size_t out_cap) {
    return encode_with_table(in, in_len, out, out_cap, b64_table, 1);
}

int litls_base64_decode(const char *in, size_t in_len, uint8_t *out, size_t *out_len) {
    return decode_with_table(in, in_len, out, out_len, b64_table);
}

size_t litls_base64url_encode(const uint8_t *in, size_t in_len, char *out, size_t out_cap) {
    return encode_with_table(in, in_len, out, out_cap, b64url_table, 0);
}

int litls_base64url_decode(const char *in, size_t in_len, uint8_t *out, size_t *out_len) {
    return decode_with_table(in, in_len, out, out_len, b64url_table);
}