{-# LANGUAGE OverloadedStrings #-}

module Data.Aeson.Helper
  ( replace
  , union
  , difference
  , pick
  ) where

import           Data.Aeson        (Value (..))
import           Data.Aeson.Key    (Key, toText)
import qualified Data.Aeson.KeyMap as KeyMap (delete, difference, insert,
                                              lookup, mapMaybeWithKey, union)
import           Data.Maybe        (catMaybes, fromMaybe)
import           Data.Text         (Text)
import qualified Data.Text         as T (isPrefixOf, stripPrefix)
import qualified Data.Vector       as V (map)

-- | Replace JSON key to a new key
--
-- >>> replace "okey" "nkey" (object [ "okey" := "value" ])
-- Object (fromList [("nkey",String "value")])
--
-- >>> replace "okey" "nkey" (String "value")
-- String "value"
replace :: Key -> Key -> Value -> Value
replace :: Key -> Key -> Value -> Value
replace Key
okey Key
nkey (Object Object
v) = Object -> Value
Object (Object -> Value) -> (Object -> Object) -> Object -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> Value -> Object -> Object
forall v. Key -> v -> KeyMap v -> KeyMap v
KeyMap.insert Key
nkey Value
ov (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Key -> Object -> Object
forall v. Key -> KeyMap v -> KeyMap v
KeyMap.delete Key
okey Object
v
  where ov :: Value
ov = Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe Value
Null (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$ Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup Key
okey Object
v

replace Key
_ Key
_ Value
v = Value
v

-- | Union two JSON
--
-- >>> union (object ["key1" .= "value1"]) (object ["key2" .= "value2"])
-- Object (fromList [("key2",String "value2"),("key1",String "value1")])
--
-- >>> union (object ["key1" .= "value1"]) (object ["key1" .= "value2"])
-- Object (fromList [("key1",String "value1")])
--
-- >>> union Null (object ["key2" .= "value2"])
-- Object (fromList [("key2",String "value2")])
--
-- >>> union (object ["key1" .= "value1"]) Null
-- Object (fromList [("key1",String "value1")])
union :: Value -> Value -> Value
union :: Value -> Value -> Value
union (Object Object
a) (Object Object
b) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Object -> Object -> Object
forall v. KeyMap v -> KeyMap v -> KeyMap v
KeyMap.union Object
a Object
b
union (Object Object
a) Value
_          = Object -> Value
Object Object
a
union Value
_ (Object Object
b)          = Object -> Value
Object Object
b
union Value
_ Value
_                   = Value
Null

-- | Difference two JSON
--
-- >>> difference  (object ["key1" .= "value1", "key2" .= "value2"]) (object ["key1" .= Null])
-- Object (fromList [("key2",String "value2")])
difference :: Value -> Value -> Value
difference :: Value -> Value -> Value
difference (Object Object
a) (Object Object
b) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Object -> Object -> Object
forall v v'. KeyMap v -> KeyMap v' -> KeyMap v
KeyMap.difference Object
a Object
b
difference (Object Object
a) Value
_          = Object -> Value
Object Object
a
difference Value
_ Value
_                   = Value
Null

-- | Pick a value from JSON
--
-- >>> pick ["key1"] $ object ["key1" .= "value1", "key2" .= "value2", "key3" .= "value3"]
-- Object (fromList [("key1",String "value1")])
--
-- >>> pick ["key1", "key2"] $ object ["key1" .= "value1", "key2" .= "value2", "key3" .= "value3"]
-- Object (fromList [("key2",String "value2"),("key1",String "value1")])
--
-- >>> pick ["key3.key4"] $ object ["key1" .= "value1", "key2" .= "value2", "key3" .= object ["key4" .= "value4"]]
-- Object (fromList [("key3",Object (fromList [("key4",String "value4")]))])
pick :: [Text] -> Value -> Value
pick :: [Text] -> Value -> Value
pick [] Value
v          = Value
v
pick [Text]
ks (Object Object
a) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ (Key -> Value -> Maybe Value) -> Object -> Object
forall v u. (Key -> v -> Maybe u) -> KeyMap v -> KeyMap u
KeyMap.mapMaybeWithKey ([Text] -> Key -> Value -> Maybe Value
doMapMaybeWithKey [Text]
ks) Object
a
pick [Text]
ks (Array Array
a)  = Array -> Value
Array (Array -> Value) -> Array -> Value
forall a b. (a -> b) -> a -> b
$ (Value -> Value) -> Array -> Array
forall a b. (a -> b) -> Vector a -> Vector b
V.map ([Text] -> Value -> Value
pick [Text]
ks) Array
a
pick [Text]
_ Value
_           = Value
Null

doMapMaybeWithKey :: [Text] -> Key -> Value -> Maybe Value
doMapMaybeWithKey :: [Text] -> Key -> Value -> Maybe Value
doMapMaybeWithKey [Text]
ks Key
key Value
v = [Text] -> Maybe Value
go [Text]
ks
  where go :: [Text] -> Maybe Value
        go :: [Text] -> Maybe Value
go [] = Maybe Value
forall a. Maybe a
Nothing
        go (Text
x:[Text]
xs)
          | Text
k Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
x = Value -> Maybe Value
forall a. a -> Maybe a
Just Value
v
          | (Text
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"." ) Text -> Text -> Bool
`T.isPrefixOf` Text
x = Value -> Maybe Value
forall a. a -> Maybe a
Just (Value -> Maybe Value) -> Value -> Maybe Value
forall a b. (a -> b) -> a -> b
$ [Text] -> Value -> Value
pick ([Maybe Text] -> [Text]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe Text] -> [Text]) -> [Maybe Text] -> [Text]
forall a b. (a -> b) -> a -> b
$ [Text] -> Text -> [Maybe Text]
nextKeys [Text]
ks Text
k) Value
v
          | Bool
otherwise = [Text] -> Maybe Value
go [Text]
xs

        k :: Text
k = Key -> Text
toText Key
key

nextKeys :: [Text] -> Text -> [Maybe Text]
nextKeys :: [Text] -> Text -> [Maybe Text]
nextKeys [] Text
_     = []
nextKeys (Text
x:[Text]
xs) Text
k = Text -> Text -> Maybe Text
T.stripPrefix (Text
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
".") Text
x Maybe Text -> [Maybe Text] -> [Maybe Text]
forall a. a -> [a] -> [a]
: [Text] -> Text -> [Maybe Text]
nextKeys [Text]
xs Text
k