module Data.Aeson.JSONPath
  ( runJSPQuery )
  where

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

import Data.Aeson.JSONPath.Parser (JSPQuery (..)
                                  , JSPSegment (..)
                                  , JSPChildSegment (..)
                                  , JSPSelector (..)
                                  , JSPWildcardT (..)
                                  , pJSPQuery)
import Protolude


runJSPQuery :: Text -> JSON.Value -> Either P.ParseError JSON.Value
runJSPQuery :: Text -> Value -> Either ParseError Value
runJSPQuery Text
query Value
document = do
  JSPQuery
jspath <- 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
<> Text -> String
forall a b. ConvertText a b => a -> b
toS Text
query) (Text -> String
forall a b. ConvertText a b => a -> b
toS Text
query)
  Value -> Either ParseError Value
forall a. a -> Either ParseError a
forall (m :: * -> *) a. Monad m => a -> m a
return (Value -> Either ParseError Value)
-> Value -> Either ParseError Value
forall a b. (a -> b) -> a -> b
$ JSPQuery -> Value -> Value
traverseJSPQuery JSPQuery
jspath Value
document


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


traverseJSPSegments :: [JSPSegment] -> JSON.Value -> JSON.Value
traverseJSPSegments :: [JSPSegment] -> Value -> Value
traverseJSPSegments [] Value
doc = Value
doc
traverseJSPSegments (JSPSegment
x:[JSPSegment]
xs) Value
doc = [JSPSegment] -> Value -> Value
traverseJSPSegments [JSPSegment]
xs (JSPSegment -> Value -> Value
traverseJSPSegment JSPSegment
x Value
doc)


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


traverseJSPChildSeg :: JSPChildSegment -> JSON.Value -> JSON.Value
traverseJSPChildSeg :: JSPChildSegment -> Value -> Value
traverseJSPChildSeg (JSPBracketed [JSPSelector]
sels) Value
doc = [JSPSelector] -> Value -> Value
traverseJSPSelectors [JSPSelector]
sels Value
doc
traverseJSPChildSeg (JSPMemberNameSH Text
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 (Text -> Key
K.fromText Text
key) Object
obj
traverseJSPChildSeg (JSPMemberNameSH Text
_) Value
_ = Value
emptyJSArray
traverseJSPChildSeg (JSPSegWildcard JSPWildcardT
JSPWildcard) Value
doc = Value
doc


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 (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f 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 Text
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 (Text -> Key
K.fromText Text
key) Object
obj
traverseJSPSelector (JSPNameSel Text
_) 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 (Maybe Value -> Array) -> Maybe Value -> Array
forall a b. (a -> b) -> a -> b
$ (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
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
1 JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
* 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
1 JSPIndexSelector -> JSPIndexSelector -> JSPIndexSelector
forall a. Num a => a -> a -> a
* 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