module Data.Aeson.JSONPath
  ( runJSPQuery
  , jsonPath)
  where

import qualified Data.Aeson                    as JSON
import qualified Data.Aeson.Key                as K
import qualified Data.Aeson.KeyMap             as KM
import qualified Data.Vector                   as V
import qualified Text.ParserCombinators.Parsec as P

import Data.Aeson.JSONPath.Parser (JSPQuery (..)
                                  , JSPSegment (..)
                                  , JSPChildSegment (..)
                                  , JSPDescSegment (..)
                                  , JSPSelector (..)
                                  , JSPWildcardT (..)
                                  , pJSPQuery)
import Data.Maybe                 (fromMaybe)
import Data.Text                  (Text)
import Language.Haskell.TH.Quote  (QuasiQuoter (..))
import Language.Haskell.TH.Syntax (lift)

import Prelude


jsonPath :: QuasiQuoter
jsonPath :: QuasiQuoter
jsonPath = QuasiQuoter
  { quoteExp :: String -> Q Exp
quoteExp = \String
query -> case Parsec String () JSPQuery
-> String -> String -> Either ParseError JSPQuery
forall s t a.
Stream s Identity t =>
Parsec s () a -> String -> s -> Either ParseError a
P.parse Parsec String () JSPQuery
pJSPQuery (String
"failed to parse query: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
query) String
query of
      Left ParseError
err -> String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Exp) -> String -> Q Exp
forall a b. (a -> b) -> a -> b
$ ParseError -> String
forall a. Show a => a -> String
show ParseError
err
      Right JSPQuery
ex -> JSPQuery -> Q Exp
forall t (m :: * -> *). (Lift t, Quote m) => t -> m Exp
forall (m :: * -> *). Quote m => JSPQuery -> m Exp
lift JSPQuery
ex
  , quotePat :: String -> Q Pat
quotePat = String -> String -> Q Pat
forall a. HasCallStack => String -> a
error String
"Error: quotePat"
  , quoteType :: String -> Q Type
quoteType = String -> String -> Q Type
forall a. HasCallStack => String -> a
error String
"Error: quoteType"
  , quoteDec :: String -> Q [Dec]
quoteDec = String -> String -> Q [Dec]
forall a. HasCallStack => String -> a
error String
"Error: quoteDec"
  }

-- | Run JSONPath query
--
-- @
-- {-\# LANGUAGE QuasiQuotes \#-}
--
-- import Data.Aeson          (Value)
-- import Data.Aeson.JSONPath (runJSPQuery, jsonPath)
--
-- book :: 'Value'
-- book = runJSPQuery [jsonPath|$.store.books[2]|] jsonDoc
-- @
runJSPQuery :: JSPQuery -> JSON.Value -> JSON.Value
runJSPQuery :: JSPQuery -> Value -> Value
runJSPQuery = JSPQuery -> Value -> Value
traverseJSPQuery


traverseJSPQuery :: JSPQuery -> JSON.Value -> JSON.Value
traverseJSPQuery :: JSPQuery -> Value -> Value
traverseJSPQuery (JSPRoot [JSPSegment]
segs) = [JSPSegment] -> Value -> Value
traverseJSPSegments [JSPSegment]
segs


traverseJSPSegments :: [JSPSegment] -> JSON.Value -> JSON.Value
traverseJSPSegments :: [JSPSegment] -> Value -> Value
traverseJSPSegments [JSPSegment]
xs Value
doc = (Value -> JSPSegment -> Value) -> Value -> [JSPSegment] -> Value
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ((JSPSegment -> Value -> Value) -> Value -> JSPSegment -> Value
forall a b c. (a -> b -> c) -> b -> a -> c
flip JSPSegment -> Value -> Value
traverseJSPSegment) Value
doc [JSPSegment]
xs


traverseJSPSegment :: JSPSegment -> JSON.Value -> JSON.Value
traverseJSPSegment :: JSPSegment -> Value -> Value
traverseJSPSegment (JSPChildSeg JSPChildSegment
jspChildSeg) Value
doc = JSPChildSegment -> Value -> Value
traverseJSPChildSeg JSPChildSegment
jspChildSeg Value
doc
traverseJSPSegment (JSPDescSeg JSPDescSegment
jspDescSeg) Value
doc = JSPDescSegment -> Value -> Value
traverseJSPDescSeg JSPDescSegment
jspDescSeg Value
doc


traverseJSPChildSeg :: JSPChildSegment -> JSON.Value -> JSON.Value
traverseJSPChildSeg :: JSPChildSegment -> Value -> Value
traverseJSPChildSeg (JSPChildBracketed [JSPSelector]
sels) Value
doc = [JSPSelector] -> Value -> Value
traverseJSPSelectors [JSPSelector]
sels Value
doc
traverseJSPChildSeg (JSPChildMemberNameSH JSPNameSelector
key) (JSON.Object Object
obj) = Value -> Maybe Value -> Value
forall a. a -> Maybe a -> a
fromMaybe Value
emptyJSArray (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$ Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KM.lookup (JSPNameSelector -> Key
K.fromText JSPNameSelector
key) Object
obj
traverseJSPChildSeg (JSPChildMemberNameSH JSPNameSelector
_) Value
_ = Value
emptyJSArray
traverseJSPChildSeg (JSPChildWildSeg JSPWildcardT
JSPWildcard) Value
doc = Value
doc


traverseJSPDescSeg :: JSPDescSegment -> JSON.Value -> JSON.Value
traverseJSPDescSeg :: JSPDescSegment -> Value -> Value
traverseJSPDescSeg (JSPDescBracketed [JSPSelector]
sels) Value
doc = Array -> Value
JSON.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 ([JSPSelector] -> Value -> Value
traverseJSPSelectors [JSPSelector]
sels) (Value -> Array
allElemsRecursive Value
doc)
traverseJSPDescSeg (JSPDescMemberNameSH JSPNameSelector
key) Value
doc = JSPNameSelector -> Value -> Value
traverseDescMembers JSPNameSelector
key Value
doc
traverseJSPDescSeg (JSPDescWildSeg JSPWildcardT
JSPWildcard) Value
doc = Array -> Value
JSON.Array (Array -> Value) -> Array -> Value
forall a b. (a -> b) -> a -> b
$ Value -> Array
allElemsRecursive Value
doc

-- TODO: Clean this super messy code, might require some refactoring
traverseDescMembers :: Text -> JSON.Value -> JSON.Value
traverseDescMembers :: JSPNameSelector -> Value -> Value
traverseDescMembers JSPNameSelector
key (JSON.Object Object
obj) = Array -> Value
JSON.Array (Array -> Value) -> Array -> Value
forall a b. (a -> b) -> a -> b
$ [Array] -> Array
forall a. [Vector a] -> Vector a
V.concat [
    Array -> (Value -> Array) -> Maybe Value -> Array
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Array
forall a. Vector a
V.empty Value -> Array
forall a. a -> Vector a
V.singleton (Maybe Value -> Array) -> Maybe Value -> Array
forall a b. (a -> b) -> a -> b
$ Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KM.lookup (JSPNameSelector -> Key
K.fromText JSPNameSelector
key) Object
obj,
    (Value -> Value) -> Array -> Array
forall a b. (a -> b) -> Vector a -> Vector b
V.map (JSPNameSelector -> Value -> Value
traverseDescMembers JSPNameSelector
key) (Value -> Array
allElemsRecursive (Array -> Value
JSON.Array (Array -> Value) -> Array -> Value
forall a b. (a -> b) -> a -> b
$ [Value] -> Array
forall a. [a] -> Vector a
V.fromList ([Value] -> Array) -> [Value] -> Array
forall a b. (a -> b) -> a -> b
$ Object -> [Value]
forall v. KeyMap v -> [v]
KM.elems Object
obj))
  ]
traverseDescMembers JSPNameSelector
key ar :: Value
ar@(JSON.Array Array
_) = Array -> Value
JSON.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 (JSPNameSelector -> Value -> Value
traverseDescMembers JSPNameSelector
key) (Value -> Array
allElemsRecursive Value
ar)
traverseDescMembers JSPNameSelector
_ Value
_ = Array -> Value
JSON.Array Array
forall a. Vector a
V.empty

traverseJSPSelectors :: [JSPSelector] -> JSON.Value -> JSON.Value
traverseJSPSelectors :: [JSPSelector] -> Value -> Value
traverseJSPSelectors [JSPSelector]
sels Value
doc = Array -> Value
JSON.Array (Array -> Value) -> Array -> Value
forall a b. (a -> b) -> a -> b
$ [Array] -> Array
forall a. [Vector a] -> Vector a
V.concat ([Array] -> Array) -> [Array] -> Array
forall a b. (a -> b) -> a -> b
$ (JSPSelector -> Array) -> [JSPSelector] -> [Array]
forall a b. (a -> b) -> [a] -> [b]
map JSPSelector -> Array
traverse' [JSPSelector]
sels
  where
    traverse' :: JSPSelector -> Array
traverse' = (JSPSelector -> Value -> Array) -> Value -> JSPSelector -> Array
forall a b c. (a -> b -> c) -> b -> a -> c
flip JSPSelector -> Value -> Array
traverseJSPSelector Value
doc

traverseJSPSelector :: JSPSelector -> JSON.Value -> V.Vector JSON.Value
traverseJSPSelector :: JSPSelector -> Value -> Array
traverseJSPSelector (JSPNameSel JSPNameSelector
key) (JSON.Object Object
obj) = Array -> (Value -> Array) -> Maybe Value -> Array
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Array
forall a. Vector a
V.empty Value -> Array
forall a. a -> Vector a
V.singleton (Maybe Value -> Array) -> Maybe Value -> Array
forall a b. (a -> b) -> a -> b
$ Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KM.lookup (JSPNameSelector -> Key
K.fromText JSPNameSelector
key) Object
obj
traverseJSPSelector (JSPNameSel JSPNameSelector
_) Value
_ = Array
forall a. Vector a
V.empty
traverseJSPSelector (JSPIndexSel JSPIndexSelector
idx) (JSON.Array Array
arr) = Array -> (Value -> Array) -> Maybe Value -> Array
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Array
forall a. Vector a
V.empty Value -> Array
forall a. a -> Vector a
V.singleton (if JSPIndexSelector
idx JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Ord a => a -> a -> Bool
>= JSPIndexSelector
0 then Array -> JSPIndexSelector -> Maybe Value
forall a. Vector a -> JSPIndexSelector -> Maybe a
(V.!?) Array
arr JSPIndexSelector
idx else Array -> JSPIndexSelector -> Maybe Value
forall a. Vector a -> JSPIndexSelector -> Maybe a
(V.!?) Array
arr (JSPIndexSelector
idx JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
+ Array -> JSPIndexSelector
forall a. Vector a -> JSPIndexSelector
V.length Array
arr))
traverseJSPSelector (JSPIndexSel JSPIndexSelector
_) Value
_ = Array
forall a. Vector a
V.empty
traverseJSPSelector (JSPSliceSel JSPSliceSelector
sliceVals) (JSON.Array Array
arr) = JSPSliceSelector -> Array -> Array
traverseJSPSliceSelector JSPSliceSelector
sliceVals Array
arr
traverseJSPSelector (JSPSliceSel JSPSliceSelector
_) Value
_ = Array
forall a. Vector a
V.empty
-- This does not work right with descendant segment, fix it later
traverseJSPSelector (JSPWildSel JSPWildcardT
JSPWildcard) Value
doc = Value -> Array
forall a. a -> Vector a
V.singleton Value
doc


traverseJSPSliceSelector :: (Maybe Int, Maybe Int, Int) -> JSON.Array -> V.Vector JSON.Value
traverseJSPSliceSelector :: JSPSliceSelector -> Array -> Array
traverseJSPSliceSelector (Maybe JSPIndexSelector
start, Maybe JSPIndexSelector
end, JSPIndexSelector
step) Array
doc = Maybe JSPIndexSelector
-> Maybe JSPIndexSelector -> JSPIndexSelector -> Array -> Array
forall {a}.
Maybe JSPIndexSelector
-> Maybe JSPIndexSelector
-> JSPIndexSelector
-> Vector a
-> Vector a
getSlice Maybe JSPIndexSelector
start Maybe JSPIndexSelector
end JSPIndexSelector
step Array
doc
  where
    -- TODO: Refactor this code to make it more pretty
    len :: JSPIndexSelector
len = Array -> JSPIndexSelector
forall a. Vector a -> JSPIndexSelector
V.length Array
doc
    normalize :: JSPIndexSelector -> JSPIndexSelector
normalize JSPIndexSelector
i = if JSPIndexSelector
i JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Ord a => a -> a -> Bool
>= JSPIndexSelector
0 then JSPIndexSelector
i else JSPIndexSelector
len JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
+ JSPIndexSelector
i

    sliceNormalized :: Vector a
-> (JSPIndexSelector, JSPIndexSelector) -> Bool -> Vector a
sliceNormalized Vector a
arr' (JSPIndexSelector
n_start, JSPIndexSelector
n_end) Bool
isStepNeg =
      let (JSPIndexSelector
lower, JSPIndexSelector
upper) = if Bool
isStepNeg then
            (JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
min (JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
max JSPIndexSelector
n_end (-JSPIndexSelector
1)) (JSPIndexSelector
lenJSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
-JSPIndexSelector
1), JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
min (JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
max JSPIndexSelector
n_start (-JSPIndexSelector
1)) (JSPIndexSelector
lenJSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
-JSPIndexSelector
1))
          else
            (JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
min (JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
max JSPIndexSelector
n_start JSPIndexSelector
0) JSPIndexSelector
len, JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
min (JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Ord a => a -> a -> a
max JSPIndexSelector
n_end JSPIndexSelector
0) JSPIndexSelector
len)
      in JSPIndexSelector -> JSPIndexSelector -> Vector a -> Vector a
forall a.
JSPIndexSelector -> JSPIndexSelector -> Vector a -> Vector a
V.slice JSPIndexSelector
lower ((if Bool
isStepNeg then JSPIndexSelector
1 else JSPIndexSelector
0)JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
+JSPIndexSelector
upperJSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
-JSPIndexSelector
lower) Vector a
arr'

    getSlice :: Maybe JSPIndexSelector
-> Maybe JSPIndexSelector
-> JSPIndexSelector
-> Vector a
-> Vector a
getSlice Maybe JSPIndexSelector
_ Maybe JSPIndexSelector
_ JSPIndexSelector
0 Vector a
_ = Vector a
forall a. Vector a
V.empty
    getSlice (Just JSPIndexSelector
st) (Just JSPIndexSelector
en) JSPIndexSelector
step' Vector a
arr =
      Vector a -> JSPIndexSelector -> Vector a
forall {a}. Vector a -> JSPIndexSelector -> Vector a
filterSlice (Vector a
-> (JSPIndexSelector, JSPIndexSelector) -> Bool -> Vector a
forall {a}.
Vector a
-> (JSPIndexSelector, JSPIndexSelector) -> Bool -> Vector a
sliceNormalized Vector a
arr (JSPIndexSelector -> JSPIndexSelector
normalize JSPIndexSelector
st, JSPIndexSelector -> JSPIndexSelector
normalize JSPIndexSelector
en) (JSPIndexSelector
step' JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Ord a => a -> a -> Bool
< JSPIndexSelector
0)) JSPIndexSelector
step'
    getSlice (Just JSPIndexSelector
st) Maybe JSPIndexSelector
Nothing JSPIndexSelector
step' Vector a
arr =
      Vector a -> JSPIndexSelector -> Vector a
forall {a}. Vector a -> JSPIndexSelector -> Vector a
filterSlice (Vector a
-> (JSPIndexSelector, JSPIndexSelector) -> Bool -> Vector a
forall {a}.
Vector a
-> (JSPIndexSelector, JSPIndexSelector) -> Bool -> Vector a
sliceNormalized Vector a
arr (JSPIndexSelector -> JSPIndexSelector
normalize JSPIndexSelector
st, JSPIndexSelector
len) (JSPIndexSelector
step' JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Ord a => a -> a -> Bool
< JSPIndexSelector
0)) JSPIndexSelector
step'
    getSlice Maybe JSPIndexSelector
Nothing (Just JSPIndexSelector
en) JSPIndexSelector
step' Vector a
arr =
      Vector a -> JSPIndexSelector -> Vector a
forall {a}. Vector a -> JSPIndexSelector -> Vector a
filterSlice (Vector a
-> (JSPIndexSelector, JSPIndexSelector) -> Bool -> Vector a
forall {a}.
Vector a
-> (JSPIndexSelector, JSPIndexSelector) -> Bool -> Vector a
sliceNormalized Vector a
arr (JSPIndexSelector
0, JSPIndexSelector -> JSPIndexSelector
normalize JSPIndexSelector
en) (JSPIndexSelector
step' JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Ord a => a -> a -> Bool
< JSPIndexSelector
0)) JSPIndexSelector
step'
    getSlice Maybe JSPIndexSelector
Nothing Maybe JSPIndexSelector
Nothing JSPIndexSelector
step' Vector a
arr = Vector a -> JSPIndexSelector -> Vector a
forall {a}. Vector a -> JSPIndexSelector -> Vector a
filterSlice Vector a
arr JSPIndexSelector
step'

    -- trying to avoid a step loop and keeping it "functional"
    filterSlice :: Vector a -> JSPIndexSelector -> Vector a
filterSlice Vector a
slice JSPIndexSelector
1   = Vector a
slice
    filterSlice Vector a
slice (-1) = Vector a -> Vector a
forall a. Vector a -> Vector a
V.reverse Vector a
slice
    filterSlice Vector a
slice JSPIndexSelector
n = if JSPIndexSelector
n JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Ord a => a -> a -> Bool
< JSPIndexSelector
0 then
          (JSPIndexSelector -> a -> Bool) -> Vector a -> Vector a
forall a. (JSPIndexSelector -> a -> Bool) -> Vector a -> Vector a
V.ifilter (\JSPIndexSelector
i a
_ -> JSPIndexSelector
i JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Integral a => a -> a -> a
`mod` (-JSPIndexSelector
n) JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Eq a => a -> a -> Bool
== JSPIndexSelector
0) (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector a
forall a. Vector a -> Vector a
V.reverse (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ JSPIndexSelector -> Vector a -> Vector a
forall a. JSPIndexSelector -> Vector a -> Vector a
V.drop (Vector a -> JSPIndexSelector
forall a. Vector a -> JSPIndexSelector
V.length Vector a
slice JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Integral a => a -> a -> a
`mod` (-JSPIndexSelector
n)) Vector a
slice
        else
          (JSPIndexSelector -> a -> Bool) -> Vector a -> Vector a
forall a. (JSPIndexSelector -> a -> Bool) -> Vector a -> Vector a
V.ifilter (\JSPIndexSelector
i a
_ -> JSPIndexSelector
i JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Integral a => a -> a -> a
`mod` JSPIndexSelector
n JSPIndexSelector -> JSPIndexSelector -> Bool
forall a. Eq a => a -> a -> Bool
== JSPIndexSelector
0) Vector a
slice

emptyJSArray :: JSON.Value
emptyJSArray :: Value
emptyJSArray = Array -> Value
JSON.Array Array
forall a. Vector a
V.empty

allElemsRecursive :: JSON.Value -> V.Vector JSON.Value
allElemsRecursive :: Value -> Array
allElemsRecursive (JSON.Object Object
obj) = [Array] -> Array
forall a. [Vector a] -> Vector a
V.concat [
    [Value] -> Array
forall a. [a] -> Vector a
V.fromList (Object -> [Value]
forall v. KeyMap v -> [v]
KM.elems Object
obj),
    [Array] -> Array
forall a. [Vector a] -> Vector a
V.concat ([Array] -> Array) -> [Array] -> Array
forall a b. (a -> b) -> a -> b
$ (Value -> Array) -> [Value] -> [Array]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Array
allElemsRecursive (Object -> [Value]
forall v. KeyMap v -> [v]
KM.elems Object
obj)
  ]
allElemsRecursive (JSON.Array Array
arr) = [Array] -> Array
forall a. [Vector a] -> Vector a
V.concat [
    Array
arr,
    [Array] -> Array
forall a. [Vector a] -> Vector a
V.concat ([Array] -> Array) -> [Array] -> Array
forall a b. (a -> b) -> a -> b
$ (Value -> Array) -> [Value] -> [Array]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Array
allElemsRecursive (Array -> [Value]
forall a. Vector a -> [a]
V.toList Array
arr)
  ]
allElemsRecursive Value
_ = Array
forall a. Vector a
V.empty