{-# LANGUAGE CPP #-}

-- | A version of aeson that parses with key order preserved.
--
-- Copyright:
--              (c) 2011-2016 Bryan O'Sullivan
--              (c) 2011 MailRank, Inc.

module Data.Aeson.Ordered
  ( Value(..)
  , Object
  , Array
  , safeUnion
  , value
  , decode
  , Data.Aeson.Ordered.toList
  , fromList
  , asObject
  , object
  , array
  , insert
  , delete
  , adjust
  , empty
  , eitherDecode
  , Data.Aeson.Ordered.lookup
  , toOrdered
  , fromOrdered
  ) where

import           Control.Applicative              hiding (empty)
import           Control.Lens                     (prism)
import           Data.Aeson.Lens                  (AsNumber (..))
import           Data.Aeson.Parser                (jstring)
import           Data.Attoparsec.ByteString       (Parser)
import           Data.Bifunctor                   (second)
import           Data.ByteString                  (ByteString)
import           Data.Data                        (Typeable)
import           Data.Hashable                    (Hashable (..))
import           Data.Scientific                  (Scientific)
import           Data.Vector                      (Vector)

import qualified Data.Aeson                       as J
import qualified Data.Attoparsec.ByteString       as A
import qualified Data.Attoparsec.ByteString.Char8 as A8
import qualified Data.ByteString.Lazy             as L
import qualified Data.HashMap.Strict              as Map
import qualified Data.HashMap.Strict.InsOrd       as OMap
import qualified Data.Text                        as T
import qualified Data.Vector                      as V

import           Hasura.Prelude                   hiding (empty, first, second)

--------------------------------------------------------------------------------
-- Copied constants from aeson

#define BACKSLASH 92
#define CLOSE_CURLY 125
#define CLOSE_SQUARE 93
#define COMMA 44
#define DOUBLE_QUOTE 34
#define OPEN_CURLY 123
#define OPEN_SQUARE 91
#define C_0 48
#define C_9 57
#define C_A 65
#define C_F 70
#define C_a 97
#define C_f 102
#define C_n 110
#define C_t 116

--------------------------------------------------------------------------------
-- Our altered type

-- | A JSON \"object\" (key\/value map). This is where this type
-- differs to the 'aeson' package.
newtype Object = Object_ { unObject_ :: InsOrdHashMap Text Value}
  deriving stock (Data, Eq, Generic, Read, Show, Typeable)
  deriving newtype (Hashable)

-- | Union the keys, ordered, in two maps, erroring on duplicates.
safeUnion :: Object -> Object -> Either String Object
safeUnion (Object_ x) (Object_ y) =
  fmap
    Object_
    (traverse
       id
       (OMap.unionWithKey
          (\k _a _b -> Left ("Duplicate key: " ++ T.unpack k))
          (fmap Right x)
          (fmap Right y)))

-- | Empty object.
empty :: Object
empty = Object_ mempty

-- | Insert before the element at index i. Think of it in terms of
-- 'splitAt', which is (take k, drop k). Deletes existing key, if any.
insert :: (Int, Text) -> Value -> Object -> Object
insert (idx, key) val =
  Object_ .
  OMap.fromList .
  uncurry (<>) .
  second ((key, val) :) .
  splitAt idx .
  OMap.toList .
  OMap.delete key .
  unObject_

-- | Lookup a key.
lookup :: Text -> Object -> Maybe Value
lookup key (Object_ omap) = OMap.lookup key omap

-- | Delete a key.
delete :: Text -> Object -> Object
delete key (Object_ omap) = Object_ (OMap.delete key omap)

adjust :: (Value -> Value) -> Text -> Object -> Object
adjust f key (Object_ omap) = Object_ (OMap.adjust f key omap)

-- | ToList a key.
toList :: Object -> [(Text,Value)]
toList (Object_ omap) = OMap.toList omap

-- | FromList a key.
fromList :: [(Text,Value)] -> Object
fromList = Object_ . OMap.fromList

-- | A JSON \"array\" (sequence).
type Array = Vector Value

-- | A JSON value represented as a Haskell value. Intentionally
-- shadowing the 'Value' from the aeson package.
data Value
  = Object !Object
  | Array !Array
  | String !Text
  | Number !Scientific
  | Bool !Bool
  | Null
  deriving stock (Data, Eq, Generic, Read, Show, Typeable)

instance Hashable Value where
  -- Lifted from Aeson's implementation for 'Value'.
  hashWithSalt s = \case
    (Object o)   ->   s `hashWithSalt` (0::Int) `hashWithSalt` o
    (Array a)    ->   foldl' hashWithSalt (s `hashWithSalt` (1::Int)) a
    (String str) ->   s `hashWithSalt` (2::Int) `hashWithSalt` str
    (Number n)   ->   s `hashWithSalt` (3::Int) `hashWithSalt` n
    (Bool b)     ->   s `hashWithSalt` (4::Int) `hashWithSalt` b
    Null         ->   s `hashWithSalt` (5::Int)

-- Adapter instance for 'lens-aeson' which lets us write optics over the numeric
-- values in ordered JSON collections as if they were plain 'Scientific' types.
instance AsNumber Value where
  _Number = prism upcast downcast
    where
      upcast = Number
      downcast v = case v of
        Number n -> Right n
        _        -> Left v

-- | Value pairs to Value
object :: [(Text, Value)] -> Value
object = Object . fromList

-- | Value list to Value
array :: [Value] -> Value
array = Array . V.fromList

-- | Convert Aeson Value to Ordered Value
toOrdered :: (J.ToJSON a) => a -> Value
toOrdered v = case J.toJSON v of
  J.Object obj    -> Object $ fromList $ map (second toOrdered) $ Map.toList obj
  J.Array arr     -> Array $ V.fromList $ map toOrdered $ V.toList arr
  J.String text   -> String text
  J.Number number -> Number number
  J.Bool boolean  -> Bool boolean
  J.Null          -> Null

-- | Convert Ordered Value to Aeson Value
fromOrdered :: Value -> J.Value
fromOrdered v = case v of
  Object obj    -> J.Object $ Map.fromList $ map (second fromOrdered) $
                   Data.Aeson.Ordered.toList obj
  Array arr     -> J.Array $ V.fromList $ map fromOrdered $ V.toList arr
  String text   -> J.String text
  Number number -> J.Number number
  Bool boolean  -> J.Bool boolean
  Null          -> J.Null

asObject :: IsString s => Value -> Either s Object
asObject = \case
  Object o -> Right o
  _        -> Left "expecting ordered object"

--------------------------------------------------------------------------------
-- Top-level entry points

eitherDecode :: L.ByteString -> Either String Value
eitherDecode = A.parseOnly value . L.toStrict

decode :: ByteString -> Maybe Value
decode = either (const Nothing) Just . A.parseOnly value

--------------------------------------------------------------------------------
-- Modified aeson parser

-- Copied from the aeson package.
arrayValues :: Parser Array
arrayValues = do
  skipSpace
  w <- A.peekWord8'
  if w == CLOSE_SQUARE
    then A.anyWord8 >> return V.empty
    else loop [] 1
  where
    loop acc !len = do
      v <- (value A.<?> "json list value") <* skipSpace
      ch <- A.satisfy (\w -> w == COMMA || w == CLOSE_SQUARE) A.<?> "',' or ']'"
      if ch == COMMA
        then skipSpace >> loop (v:acc) (len+1)
        else return (V.reverse (V.fromListN len (v:acc)))
{-# INLINE arrayValues #-}

-- Copied from aeson package.
objectValues :: Parser (InsOrdHashMap Text Value)
objectValues = do
  skipSpace
  w <- A.peekWord8'
  if w == CLOSE_CURLY
    then A.anyWord8 >> return OMap.empty
    else loop OMap.empty
 where
  -- Why use acc pattern here, you may ask? because 'H.fromList' use 'unsafeInsert'
  -- and it's much faster because it's doing in place update to the 'HashMap'!
  loop acc = do
    k <- (jstring A.<?> "object key") <* skipSpace <* (A8.char ':' A.<?> "':'")
    v <- (value A.<?> "object value") <* skipSpace
    ch <- A.satisfy (\w -> w == COMMA || w == CLOSE_CURLY) A.<?> "',' or '}'"
    let acc' = OMap.insert k v acc
    if ch == COMMA
      then skipSpace >> loop acc'
      else pure acc'
{-# INLINE objectValues #-}

-- Copied from aeson package.
value :: Parser Value
value = do
  skipSpace
  w <- A.peekWord8'
  case w of
    DOUBLE_QUOTE  -> String <$> jstring
    OPEN_CURLY    -> A.anyWord8 *> (Object . Object_ <$> objectValues)
    OPEN_SQUARE   -> A.anyWord8 *> (Array <$> arrayValues)
    C_f           -> A8.string "false" $> Bool False
    C_t           -> A8.string "true" $> Bool True
    C_n           -> A8.string "null" $> Null
    _              | w >= 48 && w <= 57 || w == 45
                  -> Number <$> A8.scientific
      | otherwise -> fail "not a valid json value"
{-# INLINE value #-}

-- Copied from aeson package.
-- | The only valid whitespace in a JSON document is space, newline,
-- carriage return, and tab.
skipSpace :: Parser ()
skipSpace = A.skipWhile $ \w -> w == 0x20 || w == 0x0a || w == 0x0d || w == 0x09
{-# INLINE skipSpace #-}