{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
-- | This module implement helper functions to read & write data
-- at bits level.
module Codec.Picture.BitWriter( BoolReader
                              , emptyBoolState
                              , BoolState
                              , byteAlignJpg
                              , getNextBitsLSBFirst
                              , getNextBitsMSBFirst 
                              , getNextBitJpg
                              , getNextIntJpg
                              , setDecodedString
                              , setDecodedStringMSB
                              , setDecodedStringJpg
                              , runBoolReader

                              , BoolWriteStateRef 
                              , newWriteStateRef
                              , finalizeBoolWriter
                              , finalizeBoolWriterGif
                              , writeBits'
                              , writeBitsGif

                              , initBoolState 
                              , initBoolStateJpg
                              , execBoolReader
                              , runBoolReaderWith
                              ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative( (<*>), (<$>) )
#endif

import Data.STRef
import Control.Monad( when )
import Control.Monad.ST( ST )
import qualified Control.Monad.Trans.State.Strict as S
import Data.Int ( Int32 )
import Data.Word( Word8, Word32 )
import Data.Bits( (.&.), (.|.), unsafeShiftR, unsafeShiftL )

import Codec.Picture.VectorByteConversion( blitVector )
import qualified Data.Vector.Storable.Mutable as M
import qualified Data.Vector.Storable as VS
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L


--------------------------------------------------
----            Reader
--------------------------------------------------
-- | Current bit index, current value, string
data BoolState = BoolState {-# UNPACK #-} !Int
                           {-# UNPACK #-} !Word8
                           !B.ByteString

emptyBoolState :: BoolState
emptyBoolState = BoolState (-1) 0 B.empty

-- | Type used to read bits
type BoolReader s a = S.StateT BoolState (ST s) a

runBoolReader :: BoolReader s a -> ST s a
runBoolReader action = S.evalStateT action $ BoolState 0 0 B.empty

runBoolReaderWith :: BoolState -> BoolReader s a -> ST s (a, BoolState)
runBoolReaderWith st action = S.runStateT action st

execBoolReader :: BoolState -> BoolReader s a -> ST s BoolState
execBoolReader st reader = S.execStateT reader st

initBoolState :: B.ByteString -> BoolState
initBoolState str = case B.uncons str of
     Nothing -> BoolState 0 0 B.empty
     Just (v, rest) -> BoolState 0 v rest

initBoolStateJpg :: B.ByteString -> BoolState
initBoolStateJpg str = 
   case B.uncons str of
     Nothing -> BoolState 0 0 B.empty
     Just (0xFF, rest) -> case B.uncons rest of
            Nothing                  -> BoolState 7 0 B.empty
            Just (0x00, afterMarker) -> BoolState 7 0xFF afterMarker
            Just (_   , afterMarker) -> initBoolStateJpg afterMarker
     Just (v, rest) -> BoolState 7 v rest

-- | Bitify a list of things to decode.
setDecodedString :: B.ByteString -> BoolReader s ()
setDecodedString str = case B.uncons str of
     Nothing        -> S.put $ BoolState      0 0 B.empty
     Just (v, rest) -> S.put $ BoolState      0 v    rest

-- | Drop all bit until the bit of indice 0, usefull to parse restart
-- marker, as they are byte aligned, but Huffman might not.
byteAlignJpg :: BoolReader s ()
byteAlignJpg = do
  BoolState idx _ chain <- S.get
  when (idx /= 7) (setDecodedStringJpg chain)

getNextBitJpg :: BoolReader s Bool
{-# INLINE getNextBitJpg #-}
getNextBitJpg = do
    BoolState idx v chain <- S.get
    let val = (v .&. (1 `unsafeShiftL` idx)) /= 0
    if idx == 0
      then setDecodedStringJpg chain
      else S.put $ BoolState (idx - 1) v chain
    return val

getNextIntJpg :: Int -> BoolReader s Int32
{-# INLINE getNextIntJpg #-}
getNextIntJpg = go 0 where
  go !acc !0 = return acc
  go !acc !n = do
    BoolState idx v chain <- S.get
    let !leftBits = 1 + fromIntegral idx
    if n >= leftBits then do
      setDecodedStringJpg chain
      let !remaining = n - leftBits
          !mask = (1 `unsafeShiftL` leftBits) - 1
          !finalV = fromIntegral v .&. mask
          !theseBits = finalV `unsafeShiftL` remaining
      go (acc .|. theseBits) remaining
    else do
      let !remaining = leftBits - n
          !mask = (1 `unsafeShiftL` n) - 1
          !finalV = fromIntegral v `unsafeShiftR` remaining
      S.put $ BoolState (fromIntegral remaining - 1) v chain
      return $ (finalV .&. mask) .|. acc


setDecodedStringMSB :: B.ByteString -> BoolReader s ()
setDecodedStringMSB str = case B.uncons str of
  Nothing        -> S.put $ BoolState      8 0 B.empty
  Just (v, rest) -> S.put $ BoolState      8 v    rest


{-# INLINE getNextBitsMSBFirst #-}
getNextBitsMSBFirst :: Int -> BoolReader s Word32
getNextBitsMSBFirst requested = go 0 requested where
  go :: Word32 -> Int -> BoolReader s Word32
  go !acc !0 = return acc
  go !acc !n = do
    BoolState idx v chain <- S.get
    let !leftBits = fromIntegral idx
    if n >= leftBits then do
      setDecodedStringMSB chain
      let !theseBits = fromIntegral v `unsafeShiftL` (n - leftBits)
      go (acc .|. theseBits) (n - leftBits)
    else do
      let !remaining = leftBits - n
          !mask = (1 `unsafeShiftL` remaining) - 1
      S.put $ BoolState (fromIntegral remaining) (v .&. mask) chain
      return $ (fromIntegral v `unsafeShiftR` remaining) .|. acc

{-# INLINE getNextBitsLSBFirst #-}
getNextBitsLSBFirst :: Int -> BoolReader s Word32
getNextBitsLSBFirst count = aux 0 count
  where aux acc 0 = return acc
        aux acc n = do
            bit <- getNextBit
            let nextVal | bit = acc .|. (1 `unsafeShiftL` (count - n))
                        | otherwise = acc
            aux nextVal (n - 1)

{-# INLINE getNextBit #-}
getNextBit :: BoolReader s Bool
getNextBit = do
    BoolState idx v chain <- S.get
    let val = (v .&. (1 `unsafeShiftL` idx)) /= 0
    if idx == 7
      then setDecodedString chain
      else S.put $ BoolState (idx + 1) v chain
    return val

-- | Bitify a list of things to decode. Handle Jpeg escape
-- code (0xFF 0x00), thus should be only used in JPEG decoding.
setDecodedStringJpg :: B.ByteString -> BoolReader s ()
setDecodedStringJpg str = case B.uncons str of
     Nothing        -> S.put $ BoolState 7 0 B.empty
     Just (0xFF, rest) -> case B.uncons rest of
            Nothing                  -> S.put $ BoolState 7 0 B.empty
            Just (0x00, afterMarker) -> -- trace "00" $ 
                S.put $ BoolState 7 0xFF afterMarker
            Just (_   , afterMarker) -> setDecodedStringJpg afterMarker
     Just (v, rest) ->
        S.put $ BoolState 7 v rest

--------------------------------------------------
----            Writer
--------------------------------------------------
defaultBufferSize :: Int
defaultBufferSize = 256 * 1024

data BoolWriteStateRef s = BoolWriteStateRef
        { bwsCurrBuffer   :: STRef s (M.MVector s Word8)
        , bwsBufferList   :: STRef s [B.ByteString]
        , bwsWrittenWords :: STRef s Int
        , bwsBitAcc       :: STRef s Word8
        , bwsBitReaded    :: STRef s Int
        }

newWriteStateRef :: ST s (BoolWriteStateRef s)
newWriteStateRef = do
    origMv <- M.new defaultBufferSize
    BoolWriteStateRef <$> newSTRef origMv
                      <*> newSTRef []
                      <*> newSTRef 0
                      <*> newSTRef 0
                      <*> newSTRef 0

finalizeBoolWriter :: BoolWriteStateRef s -> ST s L.ByteString
finalizeBoolWriter st = do
    flushLeftBits' st
    forceBufferFlushing' st
    L.fromChunks <$> readSTRef (bwsBufferList st)

forceBufferFlushing' :: BoolWriteStateRef s -> ST s ()
forceBufferFlushing' (BoolWriteStateRef { bwsCurrBuffer = vecRef
                                        , bwsWrittenWords = countRef
                                        , bwsBufferList = lstRef
                                        }) = do
    vec <- readSTRef vecRef
    count <- readSTRef countRef
    lst <- readSTRef lstRef

    nmv <- M.new defaultBufferSize
    str <- byteStringFromVector vec count

    writeSTRef vecRef nmv
    writeSTRef lstRef $ lst ++ [str]
    writeSTRef countRef 0

flushCurrentBuffer' :: BoolWriteStateRef s -> ST s ()
flushCurrentBuffer' st = do
    count <- readSTRef $ bwsWrittenWords st
    when (count >= defaultBufferSize)
         (forceBufferFlushing' st)

byteStringFromVector :: M.MVector s Word8 -> Int -> ST s B.ByteString
byteStringFromVector vec size = do
    frozen <- VS.unsafeFreeze vec
    return $ blitVector frozen 0 size

setBitCount' :: BoolWriteStateRef s -> Word8 -> Int -> ST s ()
{-# INLINE setBitCount' #-}
setBitCount' st acc count = do
    writeSTRef (bwsBitAcc st) acc
    writeSTRef (bwsBitReaded st) count

resetBitCount' :: BoolWriteStateRef s -> ST s ()
{-# INLINE resetBitCount' #-}
resetBitCount' st = setBitCount' st 0 0

pushByte' :: BoolWriteStateRef s -> Word8 -> ST s ()
{-# INLINE pushByte' #-}
pushByte' st v = do
    flushCurrentBuffer' st
    idx <- readSTRef (bwsWrittenWords st)
    vec <- readSTRef (bwsCurrBuffer st)
    M.write vec idx v
    writeSTRef (bwsWrittenWords st) $ idx + 1

flushLeftBits' :: BoolWriteStateRef s -> ST s ()
flushLeftBits' st = do
    currCount <- readSTRef $ bwsBitReaded st
    when (currCount > 0) $ do
      currWord <- readSTRef $ bwsBitAcc st
      pushByte' st $ currWord `unsafeShiftL` (8 - currCount)

-- | Append some data bits to a Put monad.
writeBits' :: BoolWriteStateRef s
           -> Word32     -- ^ The real data to be stored. Actual data should be in the LSB
           -> Int        -- ^ Number of bit to write from 1 to 32
           -> ST s ()
{-# INLINE writeBits' #-}
writeBits' st d c = do
    currWord <- readSTRef $ bwsBitAcc st
    currCount <- readSTRef $  bwsBitReaded st
    serialize d c currWord currCount
  where dumpByte 0xFF = pushByte' st 0xFF >> pushByte' st 0x00
        dumpByte    i = pushByte' st i

        serialize bitData bitCount currentWord count
            | bitCount + count == 8 = do
                     resetBitCount' st
                     dumpByte (fromIntegral $ (currentWord `unsafeShiftL` bitCount) .|.
                                                fromIntegral cleanData)

            | bitCount + count < 8 =
                let newVal = currentWord `unsafeShiftL` bitCount
                in setBitCount' st (newVal .|. fromIntegral cleanData) $ count + bitCount

            | otherwise =
                let leftBitCount = 8 - count :: Int
                    highPart = cleanData `unsafeShiftR` (bitCount - leftBitCount) :: Word32
                    prevPart = fromIntegral currentWord `unsafeShiftL` leftBitCount :: Word32

                    nextMask = (1 `unsafeShiftL` (bitCount - leftBitCount)) - 1 :: Word32
                    newData = cleanData .&. nextMask :: Word32
                    newCount = bitCount - leftBitCount :: Int

                    toWrite = fromIntegral $ prevPart .|. highPart :: Word8
                in dumpByte toWrite >> serialize newData newCount 0 0

              where cleanMask = (1 `unsafeShiftL` bitCount) - 1 :: Word32
                    cleanData = bitData .&. cleanMask     :: Word32

-- | Append some data bits to a Put monad.
writeBitsGif :: BoolWriteStateRef s
             -> Word32     -- ^ The real data to be stored. Actual data should be in the LSB
             -> Int        -- ^ Number of bit to write from 1 to 32
             -> ST s ()
{-# INLINE writeBitsGif #-}
writeBitsGif st d c = do
    currWord <- readSTRef $ bwsBitAcc st
    currCount <- readSTRef $  bwsBitReaded st
    serialize d c currWord currCount
  where dumpByte = pushByte' st

        serialize bitData bitCount currentWord count
            | bitCount + count == 8 = do
                     resetBitCount' st
                     dumpByte (fromIntegral $ currentWord  .|.
                                                (fromIntegral cleanData `unsafeShiftL` count))

            | bitCount + count < 8 =
                let newVal = fromIntegral cleanData `unsafeShiftL` count
                in setBitCount' st (newVal .|. currentWord) $ count + bitCount

            | otherwise =
                let leftBitCount = 8 - count :: Int
                    newData = cleanData `unsafeShiftR` leftBitCount :: Word32
                    newCount = bitCount - leftBitCount :: Int
                    toWrite = fromIntegral $ fromIntegral currentWord 
                                            .|. (cleanData `unsafeShiftL` count) :: Word8
                in dumpByte toWrite >> serialize newData newCount 0 0

              where cleanMask = (1 `unsafeShiftL` bitCount) - 1 :: Word32
                    cleanData = bitData .&. cleanMask     :: Word32

finalizeBoolWriterGif :: BoolWriteStateRef s -> ST s L.ByteString
finalizeBoolWriterGif st = do
    flushLeftBitsGif st
    forceBufferFlushing' st
    L.fromChunks <$> readSTRef (bwsBufferList st)

flushLeftBitsGif :: BoolWriteStateRef s -> ST s ()
flushLeftBitsGif st = do
    currCount <- readSTRef $ bwsBitReaded st
    when (currCount > 0) $ do
      currWord <- readSTRef $ bwsBitAcc st
      pushByte' st currWord

{-# ANN module "HLint: ignore Reduce duplication" #-}