module Network.DNS.Packer
    ( Packer

    , pack
    , pack'

    , getState
    , setState
    , modifyState

    , bytesWrote
    , withOffset

    , putByteString
    , putLazyByteString

    , putWord8
    , putWord16be
    , putWord32be

    , putBinary
    , liftToBinary
    )
    where

import qualified Data.Binary as Binary
import qualified Data.Binary.Put as Bin
import           Data.Bits
import qualified Data.ByteString as Strict
import qualified Data.ByteString.Lazy as Lazy
import           Data.Int
import           Data.Word


data PackingState s
    = PackingState {
        stResult     :: !Lazy.ByteString
      , stBytesWrote :: !Int64
      , stUserState  :: s
      }

newtype Packer s a = P { unP :: PackingState s -> (a, PackingState s) }

instance Monad (Packer s) where
    return a = P (\ s -> (a, s))
    m >>= k  = P (\ s -> let (a, s') = unP m s
                         in
                           unP (k a) s')
    fail err = do bytes <- get stBytesWrote
                  P (error (err
                            ++ ". Failed packing at byte position "
                            ++ show bytes))

get :: (PackingState s -> a) -> Packer s a
get f = P (\ s -> (f s, s))

set :: (PackingState s -> PackingState s) -> Packer s ()
set f = P (\ s -> ((), f s))

mkState :: Lazy.ByteString -> Int64 -> s -> PackingState s
mkState xs n s
    = PackingState {
        stResult     = xs
      , stBytesWrote = n
      , stUserState  = s
      }

pack' :: Packer s a -> s -> (Lazy.ByteString, s, a)
pack' m s
    = let (a, s') = unP m (mkState Lazy.empty 0 s)
      in
        (stResult s', stUserState s', a)

pack :: Packer s a -> s -> Lazy.ByteString
pack = (fst' .) . pack'
    where
      fst' (xs, _, _) = xs

getState :: Packer s s
getState = get stUserState

setState :: s -> Packer s ()
setState = modifyState . const

modifyState :: (s -> s) -> Packer s ()
modifyState f
    = set $ \ st -> st { stUserState = f (stUserState st) }

bytesWrote :: Integral i => Packer s i
bytesWrote = get stBytesWrote >>= return . fromIntegral

withOffset :: Int64 -> Packer s a -> Packer s a
withOffset n m
    = P $ \ s -> let (taken, dropped) = Lazy.splitAt n (stResult s)
                     padded           = Lazy.take n (taken `Lazy.append` Lazy.repeat 0)
                     tempState        = s {
                                          stResult     = padded
                                        , stBytesWrote = stBytesWrote s - Lazy.length dropped
                                        }
                     (a, tempState')  = unP m tempState
                     newState         = tempState {
                                          stResult     = replaceHead (stResult s) (stResult tempState')
                                        , stBytesWrote = max (stBytesWrote s) (stBytesWrote tempState')
                                        }
                 in
                   (a, newState)
      where
        replaceHead :: Lazy.ByteString -> Lazy.ByteString -> Lazy.ByteString
        replaceHead world newHead
            = let rest = Lazy.drop (Lazy.length newHead) world
              in
                newHead `Lazy.append` rest


putByteString :: Strict.ByteString -> Packer s ()
putByteString = putLazyByteString . Lazy.fromChunks . (:[])

putLazyByteString :: Lazy.ByteString -> Packer s ()
putLazyByteString xs
    = set $ \ st -> st {
                      stResult     = stResult st `Lazy.append` xs
                    , stBytesWrote = stBytesWrote st + Lazy.length xs
                    }

putWord8 :: Word8 -> Packer s ()
putWord8 w
    = set $ \ st -> st {
                      stResult     = stResult st `Lazy.snoc` w
                    , stBytesWrote = stBytesWrote st + 1
                    }

putWord16be :: Word16 -> Packer s ()
putWord16be w
    = do putWord8 $ fromIntegral $ (w `shiftR`  8) .&. 0xFF
         putWord8 $ fromIntegral $  w              .&. 0xFF

putWord32be :: Word32 -> Packer s ()
putWord32be w
    = do putWord8 $ fromIntegral $ (w `shiftR` 24) .&. 0xFF
         putWord8 $ fromIntegral $ (w `shiftR` 16) .&. 0xFF
         putWord8 $ fromIntegral $ (w `shiftR`  8) .&. 0xFF
         putWord8 $ fromIntegral $  w              .&. 0xFF


putBinary :: Binary.Binary a => a -> Packer s ()
putBinary = putLazyByteString . Binary.encode


liftToBinary :: s -> Packer s a -> Bin.PutM a
liftToBinary s m
    = do let (a, s') = unP m (mkState Lazy.empty 0 s)

         Bin.putLazyByteString (stResult s')
         return a