plum

module std
import std/Str
import std/Buffer
import std/Option
import std/Bool
import std/Number

# The Base64 package contains support for doing Base64 binary-to-text encodings.
#
# Implemented on top of `[]Byte`/`Buffer` (`libs/std/bytes.plum`/`buffer.plum`)
# — the original note here said a real bit-twiddling implementation couldn't
# be expressed without indexed byte-buffer mutation, which didn't exist in
# this language yet at the time it was written; it does now.
#
# `data`/the decoded payload are represented as `Str` throughout (matching
# the original public shape) rather than `[]Byte` — `Str` is itself just a
# byte array here (see `libs/std/str.plum`), so this is a plain byte string,
# not a claim that the content is valid UTF-8 text.

# Encode for PEM (RFC 1421): 64-character lines.
fun encodePEM(data: Str) -> Str =
  return encode(data, "+", "/", "=", 64, "\r\n")

# Encode for MIME (RFC 2045): 76-character lines.
fun encodeMIME(data: Str) -> Str =
  return encode(data, "+", "/", "=", 76, "\r\n")

# Encode for URLs (RFC 4648). Padding characters are stripped by default.
fun encodeURL(data: Str, pad: Bool = False) -> Str =
  return encode(data, "-", "_", pad ? "=" : "", 0, "")

# Configurable encoding. The defaults are for RFC 4648. `pad` is the padding
# character, or `""` to omit padding entirely; `linelen` is the output line
# length to wrap at (`0` disables wrapping), separated by `linesep`.
fun encode(data: Str, at62: Str = "+", at63: Str = "/", pad: Str = "=", linelen: Int = 0, linesep: Str = "\r\n") -> Str =
  at62b := at62.byteAt(0)
  at63b := at63.byteAt(0)
  has_pad := pad.length() > 0
  padb := has_pad ? pad.byteAt(0) : 0
  res := Buffer()
  n := data.length()
  i := 0
  line_chars := 0
  while i < n
    b0 := data.byteAt(i)
    has1 := i + 1 < n
    has2 := i + 2 < n
    b1 := has1 ? data.byteAt(i + 1) : 0
    b2 := has2 ? data.byteAt(i + 2) : 0
    res.writeByte(Byte(encodeByte(b0 >> 2, at62b, at63b)))
    res.writeByte(Byte(encodeByte({b0 & 0x03} << 4 | b1 >> 4, at62b, at63b)))
    line_chars = line_chars + 2
    if has1
      res.writeByte(Byte(encodeByte({b1 & 0x0F} << 2 | b2 >> 6, at62b, at63b)))
      line_chars = line_chars + 1
    else if has_pad
      res.writeByte(Byte(padb))
      line_chars = line_chars + 1
    if has2
      res.writeByte(Byte(encodeByte(b2 & 0x3F, at62b, at63b)))
      line_chars = line_chars + 1
    else if has_pad
      res.writeByte(Byte(padb))
      line_chars = line_chars + 1
    i = i + 3
    if linelen > 0 && line_chars >= linelen && i < n
      res.write(linesep)
      line_chars = 0
  return res.toStr()

# Encodes a single 6-bit value (0-63) as its base64 alphabet character.
fun encodeByte(i: Int, at62: Int, at63: Int) -> Int =
  if i < 26
    return 65 + i # 'A'..'Z'
  if i < 52
    return 97 + {i - 26} # 'a'..'z'
  if i < 62
    return 48 + {i - 52} # '0'..'9'
  if i == 62
    return at62
  return at63

# The inverse of `encodeByte`: a base64 alphabet character's 6-bit value, or
# -1 if `b` isn't one of the 64 alphabet characters (whitespace/padding are
# handled by the caller, not here).
fun decodeByte(b: Int, at62: Int, at63: Int) -> Int =
  if b >= 65 && b <= 90
    return b - 65
  if b >= 97 && b <= 122
    return b - 97 + 26
  if b >= 48 && b <= 57
    return b - 48 + 52
  if b == at62
    return 62
  if b == at63
    return 63
  return -1

fun isBase64Whitespace(b: Int) -> Bool =
  b == 32 || b == 9 || b == 10 || b == 13

# Decode for URLs (RFC 4648).
fun decodeUrl(data: Str) -> Option[Str] =
  return decode(data, "-", "_", "=")

# Configurable decoding. The defaults are for RFC 4648. Missing padding is
# not an error. Non-base64 data, other than whitespace (which can appear at
# any point), is an error (`None`).
fun decode(data: Str, at62: Str = "+", at63: Str = "/", pad: Str = "=") -> Option[Str] =
  at62b := at62.byteAt(0)
  at63b := at63.byteAt(0)
  has_pad := pad.length() > 0
  padb := has_pad ? pad.byteAt(0) : -1
  res := Buffer()
  n := data.length()
  i := 0
  bits := 0
  nbits := 0
  while i < n
    b := data.byteAt(i)
    i = i + 1
    if isBase64Whitespace(b) || b == padb
      continue
    v := decodeByte(b, at62b, at63b)
    if v < 0
      return None
    bits = {bits << 6} | v
    nbits = nbits + 6
    if nbits >= 8
      nbits = nbits - 8
      res.writeByte(Byte({bits >> nbits} & 0xFF))
  return Some(res.toStr())