module Text.JSON.JPath (jPath, jPath') where

import qualified Data.Map as Map
import Data.Maybe
import Text.RJson
import Text.ParserCombinators.Parsec.Combinator
import Text.ParserCombinators.Parsec.Char
import Text.ParserCombinators.Parsec.Prim

data Element = ObjectLookup String | ArrayLookup Int | WildcardLookup | DeepLookup deriving (Show)

-- |Evaluates JPath query on JSON String
jPath :: String -- ^ JPath query
	-> String 	-- ^ JSON as String
	-> Either String [JsonData] -- ^ Either error text or list of results
jPath query s = let json = parseJsonString s
	in either (Left) (Right . jPath' query) json

-- |Evaluates JPath query on pre-parsed JSON
jPath' :: String  -- ^ JPath query
	-> JsonData -- ^ Parsed JSON
	-> [JsonData] -- ^ List of results
jPath' query v = let parsedQuery = parseExpression query
	in either (const []) (\q -> jPathP q v) parsedQuery

expression = do
	result <- element `sepBy` slash
	eof
	return $ concat result

slash = string "/"

element :: GenParser Char st [Element]
element = do
	parsedName <- optionMaybe (deepLookup <|> wildcard <|> name)
	parsedIndex <- optionMaybe index
	return $ catMaybes [parsedName, parsedIndex]

name = do
	parsedName <- many1 (noneOf "/[]")
	return $ ObjectLookup parsedName

deepLookup = do
	string "**"
	return DeepLookup

wildcard = do
	string "*"
	return WildcardLookup

integer = do
	minus <- optionMaybe (string "-")
	number <- many1 digit
	return $ (fromMaybe "" minus) ++ number

index = do
	result <- between (string "[") (string "]") integer
	return $ ArrayLookup $ read result

parseExpression = parse expression "JPath query"

jPathP :: [Element] -> JsonData -> [JsonData]
jPathP [] v = [v]
jPathP (e:es) v = case e of
	ObjectLookup s -> case v of
		JDObject wtf -> maybe [] (jPathP es) $ s `Map.lookup` wtf
		otherwise -> []
	ArrayLookup i -> case v of
		JDArray vs -> if i >= length vs || i < 0 - length vs then
			[]
			else
			jPathP es $ vs !! (if i < 0 then length vs - abs i else i)
		otherwise -> []
	WildcardLookup -> case v of
		JDObject wtf -> concat $ map (jPathP es) (Map.elems wtf)
		JDArray vs -> concat $ map (jPathP es) vs
		otherwise -> []
	DeepLookup -> concat [jPathP (WildcardLookup:es) v, jPathP ([WildcardLookup, DeepLookup] ++ es) v]