{-# LANGUAGE CPP #-}

module PtrPoker.Poke where

import qualified Data.Text.Array as TextArray
import qualified Data.Text.Internal as TextInternal
import qualified PtrPoker.Ffi as Ffi
import qualified PtrPoker.IO.ByteString as ByteStringIO
import qualified PtrPoker.IO.Prim as PrimIO
import PtrPoker.Prelude hiding (concat)
import qualified PtrPoker.Text as Text

{-# RULES
"foldMap" forall f foldable.
  foldMap f foldable =
    Poke $ \p -> foldM (\p (Poke poker) -> poker p) p foldable
  #-}

-- |
-- Abstraction over an IO action,
-- which takes a pointer, populates it and
-- produces a pointer right after the populated data.
newtype Poke = Poke {pokePtr :: Ptr Word8 -> IO (Ptr Word8)}

instance Semigroup Poke where
  {-# INLINE [1] (<>) #-}
  Poke lIO <> Poke rIO =
    Poke (lIO >=> rIO)
  sconcat =
    concat

instance Monoid Poke where
  {-# INLINE [1] mempty #-}
  mempty =
    Poke return
  mconcat =
    concat

-- |
-- Reuses the IsString instance of 'ByteString'.
instance IsString Poke where
  fromString = byteString . fromString

-- |
-- Concatenate a foldable of pokes.
{-# INLINE [1] concat #-}
concat :: Foldable f => f Poke -> Poke
concat pokers =
  Poke (\p -> foldM (\p (Poke io) -> io p) p pokers)

-- |
-- Efficiently copy the contents of ByteString using @memcpy@.
{-# INLINE byteString #-}
byteString :: ByteString -> Poke
byteString bs =
  Poke $ \ptr -> inline ByteStringIO.pokeByteString ptr bs

-- |
-- Encode Word8 as byte, incrementing the pointer by 1.
{-# INLINE [1] word8 #-}
word8 :: Word8 -> Poke
word8 a =
  Poke (\p -> PrimIO.pokeWord8 p a $> plusPtr p 1)

-- |
-- Encode Word16 in Little-endian.
{-# INLINE [1] lWord16 #-}
lWord16 :: Word16 -> Poke
lWord16 a =
  Poke (\p -> PrimIO.pokeLEWord16 p a $> plusPtr p 2)

-- |
-- Encode Word16 in Big-endian.
{-# INLINE [1] bWord16 #-}
bWord16 :: Word16 -> Poke
bWord16 a =
  Poke (\p -> PrimIO.pokeBEWord16 p a $> plusPtr p 2)

-- |
-- Encode Word32 in Little-endian.
{-# INLINE [1] lWord32 #-}
lWord32 :: Word32 -> Poke
lWord32 a =
  Poke (\p -> PrimIO.pokeLEWord32 p a $> plusPtr p 4)

-- |
-- Encode Word32 in Big-endian.
{-# INLINE [1] bWord32 #-}
bWord32 :: Word32 -> Poke
bWord32 a =
  Poke (\p -> PrimIO.pokeBEWord32 p a $> plusPtr p 4)

-- |
-- Encode Word64 in Little-endian.
{-# INLINE [1] lWord64 #-}
lWord64 :: Word64 -> Poke
lWord64 a =
  Poke (\p -> PrimIO.pokeLEWord64 p a $> plusPtr p 8)

-- |
-- Encode Word64 in Big-endian.
{-# INLINE [1] bWord64 #-}
bWord64 :: Word64 -> Poke
bWord64 a =
  Poke (\p -> PrimIO.pokeBEWord64 p a $> plusPtr p 8)

-- |
-- Encode Int16 in Little-endian.
{-# INLINE lInt16 #-}
lInt16 :: Int16 -> Poke
lInt16 = lWord16 . fromIntegral

-- |
-- Encode Int16 in Big-endian.
{-# INLINE bInt16 #-}
bInt16 :: Int16 -> Poke
bInt16 = bWord16 . fromIntegral

-- |
-- Encode Int32 in Little-endian.
{-# INLINE lInt32 #-}
lInt32 :: Int32 -> Poke
lInt32 = lWord32 . fromIntegral

-- |
-- Encode Int32 in Big-endian.
{-# INLINE bInt32 #-}
bInt32 :: Int32 -> Poke
bInt32 = bWord32 . fromIntegral

-- |
-- Encode Int64 in Little-endian.
{-# INLINE lInt64 #-}
lInt64 :: Int64 -> Poke
lInt64 = lWord64 . fromIntegral

-- |
-- Encode Int64 in Big-endian.
{-# INLINE bInt64 #-}
bInt64 :: Int64 -> Poke
bInt64 = bWord64 . fromIntegral

-- |
-- Encode Text in UTF8.
{-# INLINE textUtf8 #-}
textUtf8 :: Text -> Poke
#if MIN_VERSION_text(2,0,0)
textUtf8 (TextInternal.Text arr off len) =
  Poke (\p -> do
    stToIO $ TextArray.copyToPointer arr off p len
    pure (plusPtr p len))
#else
textUtf8 = Text.destruct $ \arr off len ->
  Poke (\p -> Ffi.encodeText p arr (fromIntegral off) (fromIntegral len))
#endif

-- * ASCII integers

-------------------------

-- |
-- Encode Int8 as a signed ASCII decimal.
{-# INLINE [1] int8AsciiDec #-}
int8AsciiDec :: Int8 -> Poke
int8AsciiDec a =
  Poke (Ffi.pokeIntInDec (fromIntegral a))

-- |
-- Encode Int16 as a signed ASCII decimal.
{-# INLINE [1] int16AsciiDec #-}
int16AsciiDec :: Int16 -> Poke
int16AsciiDec a =
  Poke (Ffi.pokeIntInDec (fromIntegral a))

-- |
-- Encode Int32 as a signed ASCII decimal.
{-# INLINE [1] int32AsciiDec #-}
int32AsciiDec :: Int32 -> Poke
int32AsciiDec a =
  Poke (Ffi.pokeIntInDec (fromIntegral a))

-- |
-- Encode Int64 as a signed ASCII decimal.
{-# INLINE [1] int64AsciiDec #-}
int64AsciiDec :: Int64 -> Poke
int64AsciiDec a =
  Poke (Ffi.pokeLongLongIntInDec (fromIntegral a))

-- |
-- Encode Int as a signed ASCII decimal.
{-# INLINE [1] intAsciiDec #-}
intAsciiDec :: Int -> Poke
intAsciiDec a =
  Poke (Ffi.pokeLongLongIntInDec (fromIntegral a))

-- |
-- Encode Word8 as an unsigned ASCII decimal.
{-# INLINE [1] word8AsciiDec #-}
word8AsciiDec :: Word8 -> Poke
word8AsciiDec a =
  Poke (Ffi.pokeUIntInDec (fromIntegral a))

-- |
-- Encode Word16 as an unsigned ASCII decimal.
{-# INLINE [1] word16AsciiDec #-}
word16AsciiDec :: Word16 -> Poke
word16AsciiDec a =
  Poke (Ffi.pokeUIntInDec (fromIntegral a))

-- |
-- Encode Word32 as an unsigned ASCII decimal.
{-# INLINE [1] word32AsciiDec #-}
word32AsciiDec :: Word32 -> Poke
word32AsciiDec a =
  Poke (Ffi.pokeUIntInDec (fromIntegral a))

-- |
-- Encode Word64 as an unsigned ASCII decimal.
{-# INLINE [1] word64AsciiDec #-}
word64AsciiDec :: Word64 -> Poke
word64AsciiDec a =
  Poke (Ffi.pokeLongLongUIntInDec (fromIntegral a))

-- |
-- Encode Word as an unsigned ASCII decimal.
{-# INLINE [1] wordAsciiDec #-}
wordAsciiDec :: Word -> Poke
wordAsciiDec a =
  Poke (Ffi.pokeLongLongUIntInDec (fromIntegral a))

-- |
-- Encode Double as a signed ASCII decimal.
{-# INLINE doubleAsciiDec #-}
doubleAsciiDec :: Double -> Poke
doubleAsciiDec a =
  Poke $ \ptr ->
    Ffi.pokeDouble a ptr
      & fmap (plusPtr ptr . fromIntegral)

-- * Low level

-------------------------

-- |
-- Having the amount of bytes to be written precomputed,
-- executes an action,
-- which fills the pointer going downward,
-- starting from the pointer that follows the chunk.
-- I.e., you have to decrement the pointer
-- before writing the first byte,
-- decrement it again before writing the second byte and so on.
{-# INLINE sizedReverse #-}
sizedReverse :: Int -> (Ptr Word8 -> IO a) -> Poke
sizedReverse size action =
  Poke $ \ptr ->
    let afterPtr =
          plusPtr ptr size
     in action afterPtr $> afterPtr