{-# LANGUAGE OverloadedStrings #-}

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

import           Data.Aeson          (Value (..))
import           Data.Maybe          (catMaybes)
import           Data.Text           (Text)
import qualified Data.Text           as T (isPrefixOf, stripPrefix)
import qualified Data.Vector         as V (map)

import           Data.HashMap.Strict (delete, insert, lookupDefault,
                                      mapMaybeWithKey)
import qualified Data.HashMap.Strict as HM (difference, union)

-- | 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 :: Text -> Text -> Value -> Value
replace :: Text -> Text -> Value -> Value
replace okey :: Text
okey nkey :: Text
nkey (Object v :: Object
v) = Object -> Value
Object (Object -> Value) -> (Object -> Object) -> Object -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Value -> Object -> Object
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
insert Text
nkey Value
ov (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Text -> Object -> Object
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
delete Text
okey Object
v
  where ov :: Value
ov = Value -> Text -> Object -> Value
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
lookupDefault Value
Null Text
okey Object
v

replace _ _ v :: 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 a :: Object
a) (Object b :: Object
b) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Object -> Object -> Object
forall k v.
(Eq k, Hashable k) =>
HashMap k v -> HashMap k v -> HashMap k v
HM.union Object
a Object
b
union (Object a :: Object
a) _          = Object -> Value
Object Object
a
union _ (Object b :: Object
b)          = Object -> Value
Object Object
b
union _ _                   = 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 a :: Object
a) (Object b :: Object
b) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Object -> Object -> Object
forall k v w.
(Eq k, Hashable k) =>
HashMap k v -> HashMap k w -> HashMap k v
HM.difference Object
a Object
b
difference (Object a :: Object
a) _          = Object -> Value
Object Object
a
difference _ _                   = 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 [] v :: Value
v          = Value
v
pick ks :: [Text]
ks (Object a :: Object
a) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ (Text -> Value -> Maybe Value) -> Object -> Object
forall k v1 v2.
(k -> v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
mapMaybeWithKey ([Text] -> Text -> Value -> Maybe Value
doMapMaybeWithKey [Text]
ks) Object
a
pick ks :: [Text]
ks (Array a :: 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 _ _           = Value
Null

doMapMaybeWithKey :: [Text] -> Text -> Value -> Maybe Value
doMapMaybeWithKey :: [Text] -> Text -> Value -> Maybe Value
doMapMaybeWithKey ks :: [Text]
ks k :: Text
k v :: 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 (x :: Text
x:xs :: [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 -> 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

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