content: {-# LANGUAGE BinaryLiterals #-}
  {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  -- | RON version of Base64 encoding
  module RON.Base64 (
      alphabet,
      decode,
      decode60,
      decode60base32,
      decode64,
      decode64base32,
      decodeLetter,
      decodeLetter4,
      encode,
      encode60,
      encode60short,
      encode64,
      encode64base32short,
      encodeLetter,
      encodeLetter4,
      isLetter,
  ) where
  
  import RON.Prelude
  
  import Data.Bits (complement, shiftL, shiftR, (.&.), (.|.))
  import Data.ByteString qualified as BS
  import Data.ByteString.Lazy qualified as BSL
  
  import RON.Util.Word (
      Word4,
      Word6 (W6),
      Word60,
      leastSignificant4,
      leastSignificant6,
      leastSignificant60,
      safeCast,
   )
  
  -- | Base64 alphabet
  alphabet :: ByteString
  alphabet = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz~"
  
  -- | Check if a character is in the Base64 alphabet.
  isLetter :: Word8 -> Bool
  isLetter c =
      c - ord0 <= 9
          || c - ordA <= 25
          || c - orda <= 25
          || c == ord_
          || c == ordõ
  
  -- | Convert a Base64 letter to a number [0-63]
  decodeLetter :: Word8 -> Maybe Word6
  decodeLetter x
      | x < ord0 = Nothing
      | x <= ord9 = Just . leastSignificant6 $ x - ord0
      | x < ordA = Nothing
      | x <= ordZ = Just . leastSignificant6 $ x - ordA + posA
      | x == ord_ = Just $ leastSignificant6 pos_
      | x < orda = Nothing
      | x <= ordz = Just . leastSignificant6 $ x - orda + posa
      | x == ordõ = Just $ leastSignificant6 posõ
      | otherwise = Nothing
  
  -- ASCII milestone codes
  ord0, ord9, ordA, ordZ, ord_, orda, ordz, ordõ :: Word8
  ord0 = fromIntegral $ ord '0'
  ord9 = fromIntegral $ ord '9'
  ordA = fromIntegral $ ord 'A'
  ordZ = fromIntegral $ ord 'Z'
  ord_ = fromIntegral $ ord '_'
  orda = fromIntegral $ ord 'a'
  ordz = fromIntegral $ ord 'z'
  ordõ = fromIntegral $ ord '~'
  
  -- Base64 milestone codes
  posA, pos_, posa, posõ :: Word8
  posA = 10
  pos_ = 36
  posa = 37
  posõ = 63
  
  -- | Convert a subset [0-F] of Base64 letters to a number [0-15]
  decodeLetter4 :: Word8 -> Maybe Word4
  decodeLetter4 x
      | x < ord0 = Nothing
      | x <= ord9 = Just . leastSignificant4 $ x - ord0
      | x < ordA = Nothing
      | x <= ordZ = Just . leastSignificant4 $ x - ordA + posA
      | otherwise = Nothing
  
  -- | Decode a blob from a Base64 string
  decode :: ByteStringL -> Maybe ByteStringL
  decode =
      fmap (BSL.pack . go . map safeCast) . traverse decodeLetter . BSL.unpack
    where
      go = \case
          [a, b] -> decode2 a b
          [a, b, c] -> decode3 a b c
          a : b : c : d : rest -> decode4 a b c d ++ go rest
          _ -> []
      decode2 a b = [(a `shiftL` 2) .|. (b `shiftR` 4)]
      decode3 a b c =
          [ (a `shiftL` 2) .|. (b `shiftR` 4)
          , ((b .&. 0b1111) `shiftL` 4) .|. (c `shiftR` 2)
          ]
      decode4 a b c d =
          [ (a `shiftL` 2) .|. (b `shiftR` 4)
          , ((b .&. 0b1111) `shiftL` 4) .|. (c `shiftR` 2)
          , ((c .&. 0b11) `shiftL` 6) .|. d
          ]
  
  -- | Decode a 60-bit number from a Base64 string
  decode60 :: ByteString -> Maybe Word60
  decode60 =
      BS.unpack
          >>> traverse (fmap safeCast . decodeLetter)
          >=> (go 10 >>> fmap leastSignificant60)
    where
      go :: Int -> [Word8] -> Maybe Word64
      go n = \case
          [] | n >= 0 -> Just 0
          [a] | n >= 1 -> Just $ decode4 a 0 0 0
          [a, b]
              | n >= 2 -> Just $ decode4 a b 0 0
          [a, b, c]
              | n >= 3 -> Just $ decode4 a b c 0
          (a : b : c : d : rest)
              | n >= 4 -> do
                  lowerPart <- go (n - 4) rest
                  pure $ decode4 a b c d .|. (lowerPart `shiftR` 24)
          _ -> Nothing -- extra input
      decode4 :: Word8 -> Word8 -> Word8 -> Word8 -> Word64
      decode4 a b c d =
          (safeCast a `shiftL` 54)
              .|. (safeCast b `shiftL` 48)
              .|. (safeCast c `shiftL` 42)
              .|. (safeCast d `shiftL` 36)
  
  -- | Decode a 60-bit number from a Base32 string
  decode60base32 :: ByteString -> Maybe Word60
  decode60base32 =
      BS.unpack
          >>> traverse (fmap safeCast . decodeLetter)
          >=> (go12 >>> fmap leastSignificant60)
    where
      go12 :: [Word8] -> Maybe Word64
      go12 letters = do
          let (letters8, letters4) = splitAt 8 letters
              w8 = decodeBase32 8 letters8
          w4 <- go4 letters4
          pure $ (w8 `shiftL` 20) .|. w4
      go4 :: [Word8] -> Maybe Word64
      go4 letters = case splitAt 4 letters of
          (letters4, []) -> pure $ decodeBase32 4 letters4
          _ -> Nothing -- extra input
      decodeBase32 :: Int -> [Word8] -> Word64
      decodeBase32 len =
          foldl' (\acc b -> (acc `shiftL` 5) .|. safeCast b) 0
              . take len
              . (++ repeat 0)
  
  -- | Decode a 64-bit number from a Base64 string
  decode64 :: ByteString -> Maybe Word64
  decode64 s = do
      (s0, s1) <- BS.uncons s
      cons64 <$> decodeLetter4 s0 <*> decode60 s1
  
  -- | Decode a 64-bit number from a Base32 string
  decode64base32 :: ByteString -> Maybe Word64
  decode64base32 s = do
      (s0, s1) <- BS.uncons s
      cons64 <$> decodeLetter4 s0 <*> decode60base32 s1
  
  cons64 :: Word4 -> Word60 -> Word64
  cons64 v w = (safeCast v `shiftL` 60) .|. safeCast w
  
  -- | Encode a blob to a Base64 string
  encode :: ByteStringL -> ByteStringL
  encode = BSL.pack . go . BSL.unpack
    where
      go = \case
          [] -> []
          [a] -> encode1 a
          [a, b] -> encode2 a b
          a : b : c : rest -> encode3 a b c ++ go rest
      encode1 a =
          map
              (encodeLetter . leastSignificant6)
              [a `shiftR` 2, (a .&. 0b11) `shiftL` 4]
      encode2 a b =
          map
              (encodeLetter . leastSignificant6)
              [ a `shiftR` 2
              , ((a .&. 0b11) `shiftL` 4) .|. (b `shiftR` 4)
              , (b .&. 0b1111) `shiftL` 2
              ]
      encode3 a b c =
          map
              (encodeLetter . leastSignificant6)
              [ a `shiftR` 2
              , ((a .&. 0b11) `shiftL` 4) .|. (b `shiftR` 4)
              , ((b .&. 0b1111) `shiftL` 2) .|. (c `shiftR` 6)
              , c .&. 0b111111
              ]
  
  -- | Convert a number from [0..63] to a single letter
  encodeLetter :: Word6 -> Word8
  encodeLetter i = alphabet `BS.index` safeCast i
  
  -- | Convert a number from [0..15] to a single letter
  encodeLetter4 :: Word4 -> Word8
  encodeLetter4 i = alphabet `BS.index` safeCast i
  
  -- | Encode a 60-bit number to a Base64 string
  encode60 :: Word60 -> ByteString
  encode60 w =
      BS.pack
          [ encodeLetter $
              leastSignificant6 (safeCast w `shiftR` (6 * i) :: Word64)
          | i <- [9, 8 .. 0]
          ]
  
  -- | Encode a 60-bit number to a Base64 string, dropping trailing zeroes
  encode60short :: Word60 -> ByteString
  encode60short v = case safeCast v :: Word64 of
      0 -> "0"
      x -> BS.pack . map (encodeLetter . leastSignificant6) $ go 9 x
    where
      go _ 0 = []
      go i w =
          (w `shiftR` (6 * i)) .&. 0b111111
              : go (i - 1) (w .&. complement (0b111111 `shiftL` (6 * i)))
  
  -- | Encode a 64-bit number to a Base32 string, dropping trailing zeroes
  encode64base32short :: Word64 -> ByteString
  encode64base32short = \case
      0 -> "0"
      x -> BS.pack . map (encodeLetter . leastSignificant5) $ go 12 x
    where
      go _ 0 = []
      go i w =
          (w `shiftR` (5 * i)) .&. 0b11111
              : go (i - 1) (w .&. complement (0b11111 `shiftL` (5 * i)))
  
      leastSignificant5 w = W6 $ fromIntegral w .&. 0b11111
  
  -- | Encode a 64-bit number to a Base64 string
  encode64 :: Word64 -> ByteString
  encode64 w =
      encodeLetter (leastSignificant6 $ w `shiftR` 60)
          `BS.cons` encode60 (leastSignificant60 w)
  