#include "base64.hh"

#include <cstddef>
#include <cstdint>
#include <optional>
#include <span>
#include <string>
#include <string_view>
#include <vector>

namespace base64 {

namespace {

std::string_view kAlphabet(
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz"
    "0123456789+/");

[[nodiscard]]
inline char encode(uint8_t c) {
  return kAlphabet[c];
}

[[nodiscard]]
inline uint8_t decode(char c) {
  if (c >= 'A' && c <= 'Z') {
    return c - 'A';
  }
  if (c >= 'a' && c <= 'z') {
    return 26 + (c - 'a');
  }
  if (c >= '0' && c <= '9') {
    return 52 + (c - '0');
  }
  if (c == '+')
    return 62;
  if (c == '/')
    return 63;
  return 255;
}

}  // namespace

std::string encode(std::span<uint8_t const> data) {
  std::string ret;
  encode(data, ret);
  return ret;
}

void encode(std::span<uint8_t const> in, std::string& out) {
  out.reserve((in.size() * 4 + 3) / 3);
  size_t const leftover = in.size() % 3;
  size_t const end = in.size() - leftover;
  size_t i = 0;
  for (; i < end; i += 3) {
    out.push_back(encode(in[i] >> 2));
    out.push_back(encode(((in[i] & 0x03) << 4) | (in[i + 1] >> 4)));
    out.push_back(encode(((in[i + 1] & 0x0f) << 2) | (in[i + 2] >> 6)));
    out.push_back(encode(in[i + 2] & 0x3f));
  }
  switch (leftover) {  // NOLINT(bugprone-switch-missing-default-case)
    case 0:
      return;
    case 1:
      out.push_back(encode(in[i] >> 2));
      out.push_back(encode((in[i] & 0x03) << 4));
      out.push_back('=');
      break;
    case 2:
      out.push_back(encode(in[i] >> 2));
      out.push_back(encode(((in[i] & 0x03) << 4) | (in[i + 1] >> 4)));
      out.push_back(encode((in[i + 1] & 0x0f) << 2));
      break;
  }
  out.push_back('=');
}

std::optional<std::vector<uint8_t>> decode(std::string_view value) {
  std::vector<uint8_t> ret;
  if (decode(value, ret))
    return ret;
  return std::nullopt;
}

bool decode(std::string_view in, std::vector<uint8_t>& out) {
  if (in.size() % 4)
    return false;

  if (in.empty())
    return true;

  size_t pad;
  if (in.back() == '=') {
    if (in[in.size() - 2] == '=') {
      pad = 2;
    } else {
      pad = 1;
    }
  } else {
    pad = 0;
  }

  size_t const end = in.size() - (pad ? 4 : 0);
  size_t i = 0;
  for (; i < end; i += 4) {
    auto v1 = decode(in[i]);
    auto v2 = decode(in[i + 1]);
    auto v3 = decode(in[i + 2]);
    auto v4 = decode(in[i + 3]);
    if (v1 == 255 || v2 == 255 || v3 == 255 || v4 == 255)
      return false;
    out.push_back((v1 << 2) | (v2 >> 4));
    out.push_back(((v2 & 0xf) << 4) | (v3 >> 2));
    out.push_back(((v3 & 0x3) << 6) | v4);
  }
  switch (pad) {  // NOLINT(bugprone-switch-missing-default-case)
    case 0:
      break;
    case 1: {
      auto v1 = decode(in[i]);
      auto v2 = decode(in[i + 1]);
      auto v3 = decode(in[i + 2]);
      if (v1 == 255 || v2 == 255 || v3 == 255)
        return false;
      out.push_back((v1 << 2) | (v2 >> 4));
      out.push_back(((v2 & 0xf) << 4) | (v3 >> 2));
      break;
    }
    case 2: {
      auto v1 = decode(in[i]);
      auto v2 = decode(in[i + 1]);
      if (v1 == 255 || v2 == 255)
        return false;
      out.push_back((v1 << 2) | (v2 >> 4));
      break;
    }
  }
  return true;
}

}  // namespace base64