{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE CPP #-}
module Thrift.Protocol.Binary
  ( serializeBinary, deserializeBinary
  , Binary
  ) where

import Control.Monad
import Thrift.Binary.Parser
import Data.Bits
import Data.ByteString (ByteString)
import Data.ByteString.Builder
import Data.ByteString.Builder.Prim as Prim
import Data.List
import Data.Proxy
import Data.Text.Encoding
import Data.Word
import qualified Thrift.Binary.Parser as P
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as B
import qualified Data.HashMap.Strict as HashMap

import Thrift.Protocol
import Thrift.Protocol.Binary.Internal

data Binary

serializeBinary :: ThriftStruct a => a -> ByteString
serializeBinary = serializeGen (Proxy :: Proxy Binary)

deserializeBinary :: ThriftStruct a => ByteString -> Either String a
deserializeBinary = deserializeGen (Proxy :: Proxy Binary)

-- Type Macros -----------------------------------------------------------------

#define TSTOP   0
#define TBOOL   2
#define TBYTE   3
#define TDOUBLE 4
#define TI16    6
#define TI32    8
#define TI64    10
#define TSTRING 11
#define TSTRUCT 12
#define TMAP    13
#define TSET    14
#define TLIST   15
#define TFLOAT  19

version1 :: Word32
version1 = 0x80010000

versionMask :: Word32
versionMask = 0xFFFF0000

-- BinaryProtocol instance -----------------------------------------------------

instance Protocol Binary where
  -- Generators for Types
  getByteType   _ = TBYTE
  getI16Type    _ = TI16
  getI32Type    _ = TI32
  getI64Type    _ = TI64
  getFloatType  _ = TFLOAT
  getDoubleType _ = TDOUBLE
  getBoolType   _ = TBOOL
  getStringType _ = TSTRING
  getListType   _ = TLIST
  getSetType    _ = TSET
  getMapType    _ = TMAP
  getStructType _ = TSTRUCT

  -- Generators for tokens
  genMsgBegin proxy name msgType seqNum =
    B.word32BE (version1 .|. fromIntegral msgType) <>
    genText proxy name <>
    B.int32BE seqNum

  genStruct _ fields = mconcat fields <> B.word8 TSTOP

  genField _ _ ty fid _ val = B.word8 ty <> B.int16BE fid <> val
  genFieldPrim _ _ ty fid _ prim val = primBounded
    (liftFixedToBounded (Prim.word8 >*< Prim.int16BE) >*< prim)
    ((ty, fid), val)

  genList _ ty build elems =
    genListBegin ty elems <> mconcat (map build elems)

  genListPrim _ ty build elems =
    genListBegin ty elems <> primMapListBounded build elems

  genMap _ kt vt _ kbuild vbuild elems =
    genMapBegin kt vt elems <> mconcat (map build elems)
    where
      build (k, v) = kbuild k <> vbuild v

  genMapPrim _ kt vt _ kbuild vbuild elems =
    genMapBegin kt vt elems <> primMapListBounded build elems
    where
      build = kbuild >*< vbuild

  -- Generators for base types
  genBytePrim _ = liftFixedToBounded Prim.int8
  genI16Prim  _ = liftFixedToBounded Prim.int16BE
  genI32Prim  _ = liftFixedToBounded Prim.int32BE
  genI64Prim  _ = liftFixedToBounded Prim.int64BE
  genFloat _ = B.floatBE
  genDouble _ = B.doubleBE
  genBoolPrim _ = liftFixedToBounded $
    (\b -> if b then 1 else 0) >$< Prim.word8
  genByteString _ s = B.int32BE len <> byteString s
    where
      len = fromIntegral $ BS.length s

  -- Parsers for tokens
  parseMsgBegin proxy = do
    versionAndType <- getWord32be
    let
      version = versionAndType .&. versionMask
      msgType = fromIntegral (versionAndType .&. 0x000000FF)
    when (version /= version1) $ fail "parseMsgBegin: invalid version"
    name <- parseText proxy
    seqNum <- getInt32be
    return $ MsgBegin name msgType seqNum

  parseFieldBegin _ _ _ = do
    ty <- anyWord8
    if ty == TSTOP
    then pure FieldEnd
    else FieldBegin ty <$> getInt16be <*> pure False

  parseList _ p = do
    _ <- anyWord8
    len <- getInt32be
    ps <- replicateM (fromIntegral len) p
    return (fromIntegral len, ps)

  parseMap _ pk pv _ = do
    _ <- anyWord8
    _ <- anyWord8
    len <- getInt32be
    replicateM (fromIntegral len) $ (,) <$> pk <*> pv

  -- Parser for base types
  parseByte   _ = fromIntegral <$> P.anyWord8
  parseI16    _ = getInt16be
  parseI32    _ = getInt32be
  parseI64    _ = getInt64be
  parseFloat  _ = binaryFloat
  parseDouble _ = binaryDouble
  parseBool   _ = binaryBool
  parseBoolF _ _ = binaryBool
  parseByteString _ = getBuffer getInt32be BS.copy
  parseText _ = getBuffer getInt32be decodeUtf8

  parseSkip _ TBYTE   _ = skipN 1
  parseSkip _ TI16    _ = skipN 2
  parseSkip _ TI32    _ = skipN 4
  parseSkip _ TI64    _ = skipN 8
  parseSkip _ TFLOAT  _ = skipN 4
  parseSkip _ TDOUBLE _ = skipN 8
  parseSkip _ TBOOL   _ = skipN 1
  parseSkip _ TSTRING _ = getInt32be >>= skipN . fromIntegral
  parseSkip proxy TLIST _ = skipList proxy
  parseSkip proxy TSET  _ = skipList proxy
  parseSkip proxy TMAP  _ = do
    k <- anyWord8
    v <- anyWord8
    len <- getInt32be
    replicateM_ (fromIntegral len) $
      (,) <$> parseSkip proxy k Nothing <*> parseSkip proxy v Nothing
  parseSkip proxy TSTRUCT _ = do
    fbegin <- parseFieldBegin proxy 0 HashMap.empty
    case fbegin of
      FieldBegin ty _ _ -> parseSkip proxy ty Nothing *>
                           parseSkip proxy TSTRUCT Nothing
      FieldEnd -> pure ()
  parseSkip _ n _ = fail $ "unrecognized type code: " ++ show n

skipList :: Proxy Binary -> Parser ()
skipList proxy = do
  ty <- anyWord8
  len <- getInt32be
  replicateM_ (fromIntegral len) (parseSkip proxy ty Nothing)

genListBegin :: Type -> [a] -> Builder
genListBegin ty elems =
  primFixed (Prim.word8 >*< Prim.word32BE) (ty, len)
  where len = genericLength elems

genMapBegin :: Type -> Type -> [a] -> Builder
genMapBegin kt vt elems =
  primFixed
    (Prim.word8 >*< Prim.word8 >*< Prim.word32BE)
    (kt, (vt, len))
  where
    len = genericLength elems

{-# INLINE binaryBool #-}
binaryBool :: Parser Bool
binaryBool = do
  byte <- P.anyWord8
  case byte of
    0 -> pure False
    1 -> pure True
    n -> fail $ "invalid boolean value: " ++ show n