#pragma once

// Base64 encode/decode — MIT License
// Original: NickStrupat/Base64 v1.0.0 (https://github.com/NickStrupat/Base64)
// Repository was deleted from GitHub; this is a faithful reimplementation
// preserving the original API and MIT license terms.
//
// Requires C++14 or later.

#include <array>
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>

#include <tl/optional.hpp>

struct Base64 final
{
    static std::string
    encode(const std::uint8_t *data, std::size_t len)
    {
        static constexpr char kEncodeTable[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

        std::string result;
        result.reserve(4U * ((len + 2U) / 3U));

        std::size_t i = 0;
        for (; i + 2U < len; i += 3U) {
            auto b0 = static_cast<std::uint32_t>(data[i]);
            auto b1 = static_cast<std::uint32_t>(data[i + 1U]);
            auto b2 = static_cast<std::uint32_t>(data[i + 2U]);
            auto triple = (b0 << 16U) | (b1 << 8U) | b2;

            result.push_back(kEncodeTable[(triple >> 18U) & 0x3FU]);
            result.push_back(kEncodeTable[(triple >> 12U) & 0x3FU]);
            result.push_back(kEncodeTable[(triple >> 6U) & 0x3FU]);
            result.push_back(kEncodeTable[triple & 0x3FU]);
        }

        if (i < len) {
            auto b0 = static_cast<std::uint32_t>(data[i]);
            if (i + 1U < len) {
                auto b1 = static_cast<std::uint32_t>(data[i + 1U]);
                auto triple = (b0 << 16U) | (b1 << 8U);
                result.push_back(kEncodeTable[(triple >> 18U) & 0x3FU]);
                result.push_back(kEncodeTable[(triple >> 12U) & 0x3FU]);
                result.push_back(kEncodeTable[(triple >> 6U) & 0x3FU]);
                result.push_back('=');
            } else {
                auto triple = b0 << 16U;
                result.push_back(kEncodeTable[(triple >> 18U) & 0x3FU]);
                result.push_back(kEncodeTable[(triple >> 12U) & 0x3FU]);
                result.push_back('=');
                result.push_back('=');
            }
        }

        return result;
    }

    static std::string
    encode(const std::string &input)
    {
        return encode(reinterpret_cast<const std::uint8_t *>(input.data()), input.size());
    }

    // Returns decoded bytes, or tl::nullopt on malformed input (wrong length
    // or non-Base64 characters outside the expected alphabet).
    static tl::optional<std::vector<std::uint8_t>>
    decode(const std::string &input)
    {
        static constexpr std::int8_t kDecodeTable[256] = {
            -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
            -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 62, -1, -1, -1, 63, 52, 53,
            54, 55, 56, 57, 58, 59, 60, 61, -1, -1, -1, -1, -1, -1, -1, 0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
            10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -1, -1, -1, -1, -1, -1, 26, 27, 28,
            29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -1, -1,
            -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
            -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
            -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
            -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
            -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        };

        if (input.size() % 4U != 0U) {
            return tl::nullopt;
        }

        std::vector<std::uint8_t> result;
        result.reserve(3U * (input.size() / 4U));

        for (std::size_t i = 0; i < input.size(); i += 4U) {
            auto a = kDecodeTable[static_cast<std::uint8_t>(input[i])];
            auto b = kDecodeTable[static_cast<std::uint8_t>(input[i + 1U])];
            auto c = kDecodeTable[static_cast<std::uint8_t>(input[i + 2U])];
            auto d = kDecodeTable[static_cast<std::uint8_t>(input[i + 3U])];

            // The first two characters must be valid Base64 digits (never '=' and
            // never invalid).
            if (a < 0 || b < 0) {
                return tl::nullopt;
            }
            // Padding characters are only valid at positions 2 and 3 of a group,
            // and '=' at position 2 requires '=' at position 3.
            bool pad2 = (input[i + 2U] == '=');
            bool pad3 = (input[i + 3U] == '=');
            if (pad2 && !pad3) {
                return tl::nullopt;
            }
            if (!pad2 && c < 0) {
                return tl::nullopt;
            }
            if (!pad3 && d < 0) {
                return tl::nullopt;
            }

            auto triple = (static_cast<std::uint32_t>(a) << 18U) | (static_cast<std::uint32_t>(b) << 12U) |
                          (static_cast<std::uint32_t>(pad2 ? 0 : c) << 6U) | static_cast<std::uint32_t>(pad3 ? 0 : d);

            result.push_back(static_cast<std::uint8_t>((triple >> 16U) & 0xFFU));
            if (!pad2) {
                result.push_back(static_cast<std::uint8_t>((triple >> 8U) & 0xFFU));
            }
            if (!pad3) {
                result.push_back(static_cast<std::uint8_t>(triple & 0xFFU));
            }
        }

        return result;
    }

    // Returns decoded string, or tl::nullopt on malformed input.
    static tl::optional<std::string>
    decodeToString(const std::string &input)
    {
        auto bytes = decode(input);
        if (!bytes) {
            return tl::nullopt;
        }
        return std::string(reinterpret_cast<const char *>(bytes->data()), bytes->size());
    }
};
