{- | I wanted to insert a description here, but got tired fighting with haddock escaping. Documentation: <http://bitcheese.net/wiki/code/hjpath>
-}

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

import qualified Data.Map as Map
import Data.List
import Data.Maybe
import Text.HJson as JSON
import Text.ParserCombinators.Parsec.Combinator
import Text.ParserCombinators.Parsec.Char
import Text.ParserCombinators.Parsec.Prim

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

-- |Evaluates JPath query on JSON String
jPath :: String -- ^ JPath query
	-> String 	-- ^ JSON as String
	-> Either String [Json] -- ^ (Left parsing error) or (Right results)
jPath query s = let json = JSON.fromString s
	in either (Left) (Right . jPath' query) json

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

-- |Modifies JSON content under JPath expression
jPathModify :: String -- ^ JPath query
	-> (Json -> Json) -- ^ Element modifier function
	-> String -- ^ JSON as string
	-> Either String String -- ^ (Left parsing error) or (Right modified JSON as string)
jPathModify query modifier json = either (Left) (Right . JSON.toString . jPathModify' query modifier) $ JSON.fromString json

-- |jPathModify for pre-parsed JSON
jPathModify' :: String -- ^ JPath query
	-> (Json -> Json) -- ^ Element modifier function
	-> Json -- ^ JSON
	-> Json -- ^ modified JSON
jPathModify' query modifier json = either (const json) (\q -> (jPathModifyP q modifier json)) $ parseExpression query

-- private functions

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] -> Json -> [Json]
jPathP [] v = [v]
jPathP (e:es) v = case e of
	ObjectLookup s -> case v of
		JObject wtf -> maybe [] (jPathP es) $ s `Map.lookup` wtf
		otherwise -> []
	ArrayLookup i -> case v of
		JArray vs -> if i >= genericLength vs || i < 0 - genericLength vs then
			[]
			else
			jPathP es $ vs `genericIndex` (if i < 0 then genericLength vs - abs i else i)
		otherwise -> []
	WildcardLookup -> case v of
		JObject wtf -> concat $ map (jPathP es) (Map.elems wtf)
		JArray vs -> concat $ map (jPathP es) vs
		otherwise -> []
	DeepLookup -> concat [jPathP (WildcardLookup:es) v, jPathP ([WildcardLookup, DeepLookup] ++ es) v]

jPathModifyP :: [Element] -> (Json -> Json) -> Json -> Json
jPathModifyP [] modifier v = modifier v
jPathModifyP (e:es) modifier v = let traverse = jPathModifyP es modifier in case e of
	ObjectLookup s -> case v of
		JObject wtf -> JObject $ Map.alter (Just . traverse . fromMaybe (JObject Map.empty)) s wtf
		otherwise -> v
	ArrayLookup i -> case v of
		JArray vs -> if i >= genericLength vs || i < 0 - genericLength vs then
			JArray vs
			else
			let num = if i < 0 then genericLength vs - abs i else i;
				(hd, tl) = genericSplitAt (num + 1) vs
				in JArray $ concat [init hd, [traverse (last hd)], tl]
				-- Ugly. why Data.List doesn't have 'modify by index' function
		otherwise -> v
	WildcardLookup -> case v of
		JObject wtf -> JObject $ Map.map (traverse) wtf
		JArray vs -> JArray $ map (traverse) vs
		otherwise -> v
	DeepLookup -> jPathModifyP (WildcardLookup:es) modifier  (jPathModifyP ([WildcardLookup, DeepLookup] ++ es) modifier v)