{-
  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 #-}
{-# LANGUAGE TypeApplications #-}

-- | Implementation of the Thrift Simple JSON Protocol
module Thrift.Protocol.JSON
  ( serializeJSON, serializePrettyJSON
  , deserializeJSON
  , JSON
  , keyToStr, hmMapKeys
  ) where

import Control.Arrow hiding (loop)
import Control.Monad
import Data.Aeson
import Data.Aeson.Encode.Pretty
import qualified Data.Aeson.Encoding as AE
import Thrift.Binary.Parser as P
import Data.ByteString (ByteString)
import Data.ByteString.Builder as B
import Data.ByteString.Builder.Prim as Prim
import Data.Function
import Data.Hashable
import Data.HashMap.Strict (HashMap)
import Data.List
import Data.Proxy
import Data.Text (Text)
import Data.Text.Encoding
import Data.Word8
import GHC.Float
import qualified Data.ByteString.Lazy as LBS
import qualified Data.HashMap.Strict as HashMap

import Thrift.Protocol
import Thrift.Protocol.JSON.Base64
import Thrift.Protocol.JSON.String

-- | Implementation of TSimpleJSONProtocol
serializeJSON :: ThriftSerializable a => a -> ByteString
serializeJSON = serializeGen (Proxy :: Proxy JSON)

-- | Pretty printed version of SimpleJSONProtocol
-- We still use Aeson here for now
serializePrettyJSON :: ThriftSerializable a => a -> ByteString
serializePrettyJSON = LBS.toStrict . encodePretty

deserializeJSON :: ThriftSerializable a => ByteString -> Either String a
deserializeJSON = deserializeGen (Proxy :: Proxy JSON)

-- Helper for converting Maps --------------------------------------------------

-- The keys of JSON Maps have to be Text, but in thrift they can be anything, so
-- we have these helpers to fix them up

keyToStr :: ToJSON k => k -> Text
keyToStr = decodeUtf8 . LBS.toStrict . encode

hmMapKeys :: (Eq b, Hashable b) => (a -> b) -> HashMap a x -> HashMap b x
hmMapKeys f = HashMap.fromList . map (first f) . HashMap.toList

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

data JSON

#define COMMA 44
#define COLON 58
#define QUOTE 34
#define LCURLY 123
#define RCURLY 125
#define LBRACK 91
#define RBRACK 93

#define C_t 116
#define C_f 102
#define C_n 110

version1 :: Int
version1 = 1

-- These Type codes don't appear in the serialized data, rather we just make
-- them up in order to make sure that the value in the serialized data matches
-- what we expect. There are several types that we can't distinguish between
-- (eg numbers, lists/sets, etc)
numberTy, boolTy, strTy, listTy, structTy, nullTy :: Word8
numberTy = 0
boolTy   = 1
strTy    = 2
listTy   = 3
structTy = 4
nullTy   = 5

instance Protocol JSON where
  -- Generators for Types
  getByteType   _ = numberTy
  getI16Type    _ = numberTy
  getI32Type    _ = numberTy
  getI64Type    _ = numberTy
  getFloatType  _ = numberTy
  getDoubleType _ = numberTy
  getBoolType   _ = boolTy
  getStringType _ = strTy
  getListType   _ = listTy
  getSetType    _ = listTy
  getMapType    _ = structTy
  getStructType _ = structTy

  -- Generators for tokens
  genMsgBegin _ name msgTy seqNum =
    "[" <> B.intDec version1   <>
    "," <> genEscString name <>
    "," <> B.word8Dec msgTy    <>
    "," <> B.int32Dec seqNum   <>
    ","
  genMsgEnd _ = "]"

  genStruct _ fields = "{" <> mconcat (intersperse "," fields) <> "}"

  genField _ name _ _ _ val = genEscString name <> ":" <> val
  genFieldPrim _ name _ _ _ prim val =
    genEscString name <> ":" <> primBounded prim val

  genList _ _ build elems =
    "[" <> mconcat (intersperse "," (map build elems)) <> "]"

  genListPrim p ty prim = genList p ty (primBounded prim)

  genMap _ _ _ isString kbuild vbuild elems =
    "{" <> mconcat (intersperse "," (map build elems)) <> "}"
    where
      build (k, v)
        | isString  = kbuild k <> ":" <> vbuild v
        | otherwise = "\"" <> kbuild k <> "\":" <> vbuild v

  genMapPrim p k v b kb vb = genMap p k v b (primBounded kb) (primBounded vb)

  -- Generators for base types
  genBytePrim _ = Prim.int8Dec
  genI16Prim  _ = Prim.int16Dec
  genI32Prim  _ = Prim.int32Dec
  genI64Prim  _ = Prim.int64Dec
  genFloat _ num
    | isNaN num = Prim.primMapListFixed Prim.char7 "\"NaN\""
    -- Generate x rather than x.0, like C++'s serializer
    | num == fromInteger (round num) = B.intDec $ round num
    | otherwise = floatDec num
  genDouble _ num
    | isNaN num = Prim.primMapListFixed Prim.char7 "\"NaN\""
    -- Generate x rather than x.0, like C++'s serializer
    | num == fromInteger (round num) = B.intDec $ round num
    | otherwise = doubleDec num
  genBoolPrim _ = (condB id `on` liftFixedToBounded)
    (const ('t', ('r', ('u', 'e'))) >$<
     (Prim.char7 >*< Prim.char7 >*< Prim.char7 >*< Prim.char7))
    (const ('f', ('a', ('l', ('s', 'e')))) >$<
     (Prim.char7 >*< Prim.char7 >*< Prim.char7 >*< Prim.char7 >*< Prim.char7))

  genText _ = genEscString
  genBytes _ bs = "\"" <> encodeBase64 bs <> "\""

  -- Parsers for tokens
  parseMsgBegin _ = do
    _       <- skipSpaces *> P.word8 LBRACK
    version <- skipSpaces *> parseJSONInt
    when (version /= version1) $ fail "parseMsgBegin: invalid version"
    _       <- skipSpaces *> P.word8 COMMA
    name    <- parseJSONString
    _       <- skipSpaces *> P.word8 COMMA
    msgType <- parseJSONInt
    _       <- skipSpaces *> P.word8 COMMA
    seqNum  <- parseJSONInt
    _       <- skipSpaces *> P.word8 COMMA
    return $ MsgBegin name msgType seqNum
  parseMsgEnd _ = void $ skipSpaces *> P.word8 RBRACK

  parseFieldBegin _ lastId idMap = do
    w <- skipSpaces *> anyWord8
    case w of
      -- If we get a '{' then the struct might be empty
      LCURLY | lastId == 0 -> do
        c <- skipSpaces *> P.peek
        case c of
          RCURLY -> FieldEnd <$ anyWord8
          _ -> go
      COMMA -> go
      RCURLY -> pure FieldEnd
      _  -> fail $ "no field " ++ show w ++ " " ++ show lastId
    where
      go = do
        name <- parseJSONString
        _    <- skipSpaces *> P.word8 COLON
        case HashMap.lookup name idMap of
          Just fid -> do
            c <- skipSpaces *> P.peek
            let
              ty = case c of
                LCURLY -> structTy
                LBRACK -> listTy
                QUOTE  -> strTy
                C_f    -> boolTy
                C_t    -> boolTy
                C_n    -> nullTy
                _      -> numberTy
            pure $ FieldBegin ty fid False
          -- Inserting a 0xFF type code ensures that this won't match with
          -- anything
          Nothing  -> pure $ FieldBegin 0xFF (-1) False

  -- Parser for base types
  parseByte   _ = parseJSONInt
  parseI16    _ = parseJSONInt
  parseI32    _ = parseJSONInt
  parseI64    _ = parseJSONInt
  parseFloat  _ = double2Float <$> parseJSONDouble
  parseDouble _ = parseJSONDouble
  parseBool   _ = parseJSONBool
  parseBoolF _ _ = parseJSONBool
  parseText _ = parseJSONString
  parseBytes _ = do
    raw <- P.word8 QUOTE *> P.takeWhile (/= QUOTE) <* P.word8 QUOTE
    return $ decodeBase64 raw

  parseList _ p = do
    _ <- skipSpaces *> P.word8 LBRACK
    w <- skipSpaces *> P.peek
    if w == RBRACK then anyWord8 *> pure (0,[]) else loop 1 []
      where
        loop !n xs = do
          x <- skipSpaces *> p
          w <- skipSpaces *> anyWord8
          case w of
            RBRACK -> pure (n, reverse (x:xs))
            COMMA  -> loop (n+1) (x:xs)
            _ -> fail "invalid list"

  parseMap _ pk pv isString = do
    _ <- skipSpaces *> P.word8 LCURLY
    w <- skipSpaces *> P.peek
    if w == RCURLY then anyWord8 *> pure [] else loop
      where
        loop = do
          skipSpaces
          key <- if isString then pk else P.word8 QUOTE *> pk <* P.word8 QUOTE
          _   <- skipSpaces *> P.word8 COLON
          val <- skipSpaces *> pv
          w <- skipSpaces *> anyWord8
          case w of
            RCURLY -> pure [(key, val)]
            COMMA  -> ((key, val) :) <$> loop
            _ -> fail "invalid map"

  parseSkip _ _ _ = do
    c <- skipSpaces *> P.peek
    case c of
      LCURLY -> skipObj
      LBRACK -> skipList
      QUOTE  -> void $ parseJSONString
      C_f    -> void $ string "false"
      C_t    -> void $ string "true"
      C_n    -> void $ string "null"
      _ | c >= _0 && c <= _9 || c == _hyphen -> void double
        | otherwise -> fail "not a valid json value"

skipObj :: Parser ()
skipObj = do
  _ <- P.word8 LCURLY
  _ <- sepBy skipInner $ skipSpaces *> P.word8 COMMA
  _ <- skipSpaces *> P.word8 RCURLY
  return ()
  where
    skipInner = do
      _ <- skipSpaces *> parseSkip (Proxy @JSON) 1 Nothing
      _ <- skipSpaces *> P.word8 COLON
      skipSpaces *> parseSkip (Proxy @JSON) 1 Nothing

skipList :: Parser ()
skipList = do
  _ <- P.word8 LBRACK
  _ <- sepBy skipInner $ skipSpaces *> P.word8 COMMA
  _ <- skipSpaces *> P.word8 RBRACK
  return ()
  where
    skipInner = skipSpaces *> parseSkip (Proxy @JSON) 1 Nothing

-- Use Aeson to escape a string
genEscString :: Text -> B.Builder
genEscString = fromEncoding . AE.text

{-# INLINE parseJSONInt #-}
parseJSONInt :: Integral a => Parser a
parseJSONInt = lexeme $ signed decimal

{-# INLINE parseJSONBool #-}
parseJSONBool :: Parser Bool
parseJSONBool = do
  skipSpaces
  w <- anyWord8
  if | w == C_t -> True <$ string "rue"
     | w == C_f -> False <$ string "alse"
     | otherwise -> fail "expected boolean"

{-# INLINE parseJSONDouble #-}
parseJSONDouble :: Parser Double
parseJSONDouble = do
  c <- skipSpaces *> P.peek
  if | c == QUOTE -> string "\"NaN\"" >> return (0/0)
     | otherwise -> double


{-# INLINE lexeme #-}
lexeme :: Parser a -> Parser a
lexeme = (skipSpaces *>)