-- Copyright 2015 Alvaro J. Genial (http://alva.ro) -- see LICENSE.md for more.
-- Informed by Mario Rodriguez's C++ implementation.

module Codec.Binary.Base91.String (decode, encode) where

import qualified Codec.Binary.Base91 as B91
import Data.Bits ((.&.), (.|.), shiftL, shiftR)
import Data.Char (ord)
import Data.List (foldl')
import Data.Word (Word8)

-- | Encodes octets ('[Word8]') to a 'String' in Base91; the opposite of 'decode'.
encode :: [Word8] -> String
encode = g . foldl' f (0, 0, []) where

  f :: (Int, Int, [Char]) -> Word8 -> (Int, Int, [Char])
  f (queue, nbits, cs) w =
    let queue' = queue .|. (fromIntegral w `shiftL` nbits)
        nbits' = nbits + 8
    in if nbits' <= 13 then (queue', nbits', cs) else
      let val  = queue' .&. 8191
          val' = queue' .&. 16383
          (v, q, n)   = if val > 88
            then (val,  queue' `shiftR` 13, nbits' - 13)
            else (val', queue' `shiftR` 14, nbits' - 14)
          trail       = [B91.encoding !! (v `mod` 91),
                         B91.encoding !! (v `div` 91)]
      in (q, n, cs ++ trail)

  g :: (Int, Int, [Char]) -> [Char]
  g (_,     0,     cs) = cs
  g (queue, nbits, cs) = cs ++ [y] ++ z
    where y = B91.encoding !! (queue `mod` 91)
          z | nbits > 7 || queue > 90 = [B91.encoding !! (queue `div` 91)]
            | otherwise               = []

-- | Decodes octets ('[Word8]') from a 'String' in Base91; the opposite of 'encode'.
decode :: String -> [Word8]
decode = g . foldl' f (0, 0, -1, []) where

  f :: (Int, Int, Int, [Word8]) -> Char -> (Int, Int, Int, [Word8])
  f (queue, nbits, val, ws) c =
    let d = fromIntegral $ B91.decoding !! ord c
     in if d   == 91 then (queue, nbits, val, ws) else
        if val == -1 then (queue, nbits, d,   ws) else
            let v = val + (d * 91)
                q = queue .|. (v `shiftL` nbits)
                n = nbits + (if (v .&. 8191) > 88 then 13 else 14)
                (queue', nbits', trail) = if n - 8 > 7
                  then (q `shiftR` 16, n - 16, [q, q `shiftR` 8])
                  else (q `shiftR` 8,  n - 8,  [q])
             in (queue', nbits', -1, ws ++ map fromIntegral trail)

  g :: (Int, Int, Int, [Word8]) -> [Word8]
  g (_,     _,     -1,  ws) = ws
  g (queue, nbits, val, ws) = ws ++ [fromIntegral $ queue .|. (val `shiftL` nbits)]