module LengthPrefixedBE
    ( LengthPrefixedBE(..)
    , encode_bigendian
    , decode_bigendian
    ) where

import qualified Data.ByteString.Lazy as L
import Data.Bits
import Data.Binary
import Data.Binary.Get
import Data.Binary.Put (putWord32be, putLazyByteString)
import Data.Int

{-
   From RFC4251...
 
   string

      Arbitrary length binary string.  Strings are allowed to contain
      arbitrary binary data, including null characters and 8-bit
      characters.  They are stored as a uint32 containing its length
      (number of bytes that follow) and zero (= empty string) or more
      bytes that are the value of the string.  Terminating null
      characters are not used.

   mpint ( LengthPrefixedBE )

      Represents multiple precision integers in two's complement format,
      stored as a string, 8 bits per byte, MSB first.  Negative numbers
      have the value 1 as the most significant bit of the first byte of
      the data partition.  If the most significant bit would be set for
      a positive number, the number MUST be preceded by a zero byte.
      Unnecessary leading bytes with the value 0 or 255 MUST NOT be
      included.  The value zero MUST be stored as a string with zero
      bytes of data.
-}

newtype LengthPrefixedBE = LengthPrefixedBE Integer

instance Binary LengthPrefixedBE where

    put (LengthPrefixedBE n) = do
        putWord32be len
        putLazyByteString bytes
     where
        bytes = encode_bigendian n
        len = fromIntegral (L.length bytes) :: Word32

    get = do
        len <- get
        bs <- getLazyByteString (word32_to_int64 len)
        return . LengthPrefixedBE $ decode_bigendian bs
     where
        word32_to_int64 :: Word32 -> Int64
        word32_to_int64 = fromIntegral
        


encode_bigendian :: (Integral a, Bits a) => a -> L.ByteString
encode_bigendian n =
        if (bit /= sbyte) 
         then sbyte `L.cons` bytes
         else bytes
     where
        bytes = L.reverse $ unroll n
        sbyte :: Word8
        sbyte = if n<0 then 0xFF else 0
        bit = if L.null bytes 
                then 0x00
                else fromIntegral ((fromIntegral (L.head bytes) :: Int8) `shiftR` 7)

        unroll :: (Integral a, Bits a) => a -> L.ByteString
        unroll = L.unfoldr step 
            -- TODO: Is reversing L.unfoldr more or less efficient
            --  than using Data.List.unfoldr ?
            --  Probably Data.ByteString.Lazy should export an unfoldrEnd
            --  function that efficiently unfolds reversed bytestrings.
          where
            step 0 = Nothing
            step (-1) = Nothing
            step i = Just (fromIntegral i, i `shiftR` 8)

decode_bigendian :: (Num a, Bits a) => L.ByteString -> a
decode_bigendian bs = if isneg then n - 256^(L.length bs)
                               else n
 where
    n = L.foldl (\a b -> a `shiftL` 8 .|. fromIntegral b) 0 bs
    isneg = not (L.null bs) &&  L.head bs .&. 0x80 /= 0