-- |  See also: <www.haskell.org/haskellwiki/HXT> 
module Data.JSON2.Query
(
-- * Data Types
     JFilter
-- * Filtering
    ,isObj, isArr
    ,isStr ,isStrBy
    ,isNum ,isNumBy
    ,isBool, isTrue, isFalse
    ,isNull,isAtomic
    ,getFromObj, getFromKey, getFromKeys
    ,getFromArr ,getFromIndex ,getFromIndexes
    ,getChildern
-- * Filter Combinators
    ,(>>>) ,(<+>)
    ,orElse ,when ,guards
    ,deep, deepObj, deepArr
)
where
import qualified Data.Map as M (empty, singleton, unionWith, elems, union, lookup)
import qualified Data.Maybe as Mb (catMaybes)
import qualified Data.List as L (nub)
import Data.JSON2

type JFilter = Json -> Jsons

infixl 1 >>>
infixr 2 <+>

-- Fitering

-- | Filter JSON objects.
isObj :: JFilter
isObj (JObject o) = [JObject o]
isObj _              = []

-- | Filter JSON arrays.
isArr :: JFilter
isArr (JArray a)  = [JArray a]
isArr _              = []

-- | Filter JSON strings.
isStr :: JFilter
isStr (JString s)  = [JString s]
isStr _            = []

-- | Predicative filter JSON strings.
isStrBy :: (String -> Bool) -> JFilter
isStrBy p (JString s)  = if p s then [JString s] else []
isStrBy _ _           = []

-- | Filter JSON numbers.
isNum :: JFilter
isNum (JNumber n) = [JNumber n]
isNum _           = []


-- | Predicative filter JSON numbers.
isNumBy :: Fractional a => (a -> Bool) -> JFilter
isNumBy p (JNumber n) = if p (fromRational n) then [JNumber n] else []
isNumBy _ _          = []

-- | Filter JSON Bool.
isBool :: JFilter
isBool (JBool p)   = [JBool p]
isBool _           = []

-- | Filter JSON True.
isTrue :: JFilter
isTrue (JBool True)   = [JBool True]
isTrue _           = []

-- | Filter JSON False.
isFalse :: JFilter
isFalse (JBool False)   = [JBool False]
isFalse _           = []

-- | Filter JSON null.
isNull :: JFilter
isNull JNull = [JNull]
isNull _ = []

-- | Filter primitive types.
isAtomic :: JFilter
isAtomic (JString s)  = [JString s]
isAtomic (JNumber n)  = [JNumber n]
isAtomic (JBool p)    = [JBool p]
isAtomic JNull        = [JNull]
isAtomic _            = []


-- | Get elements from object with key.
getFromKey :: String -> JFilter
getFromKey k (JObject m) = Mb.catMaybes  [(M.lookup k m)]
getFromKey _ _                = []

-- | Get elements from object with keys.
--  
--   Examles:
--
-- > citys :: Jsons
-- > citys =
-- >   [
-- >     "Europa"    ==> "Ukraine" ==> ["Kiyv", "Gitomir", "Lviv"],
-- >     "Asia"      ==> "Japan"   ==> ["Tokyo"],
-- >     "Europa"    ==> "UK"      ==> ["London", "Glasgow"],
-- >     "Europa"    ==> "Germany" ==> ["Berlin", "Bonn"],
-- >     "America"   ==> "USA"     ==> ["NewYork"],
-- >     "America"   ==> "Canada"  ==> ["Toronto"],
-- >     "Australia" ==> ["Melburn", "Adelaida"]
-- >   ]
-- > jCitys = foldl unionRecObj emptyObj citys
-- > ex1 = pprint jCitys
-- > {
-- >   "America": {
-- >     "Canada": ["Toronto"],
-- >     "USA": ["NewYork"]
-- >   },
-- >   "Asia": {
-- >     "Japan": ["Tokyo"]
-- >   },
-- >   "Australia": ["Melburn", "Adelaida"],
-- >   "Europa": {
-- >     "Germany": ["Berlin", "Bonn"],
-- >     "UK": ["London", "Glasgow"],
-- >     "Ukraine": ["Kiyv", "Gitomir", "Lviv"]
-- >   }
-- > }
-- > query2 = getFromKeys ["Europa", "America", "Atlantida"]
-- > ex2 = pprints $ query2 jCitys
-- > [
-- >   {
-- >     "Germany": ["Berlin", "Bonn"],
-- >     "UK": ["London", "Glasgow"],
-- >     "Ukraine": ["Kiyv", "Gitomir", "Lviv"]
-- >   },{
-- >     "Canada": ["Toronto"],
-- >     "USA": ["NewYork"]
-- >   }
-- > ]
getFromKeys :: [String] -> JFilter
getFromKeys ks (JObject m) = Mb.catMaybes $ map (\k -> (M.lookup k m)) (L.nub ks) 
getFromKeys _ _            = []

-- | Get all elements from object.
getFromObj :: JFilter
getFromObj (JObject o) = M.elems o
getFromObj _              = []

-- | Get all elements from array.
getFromArr :: JFilter
getFromArr (JArray a)  = a
getFromArr _           = []

-- | Get element from array with index.
getFromIndex :: Int -> JFilter
getFromIndex i (JArray a) = if i < length a  then [a !! i] else []
getFromIndex _ _          = []

-- | Get elements from array with indexes.
getFromIndexes :: [Int] -> JFilter
getFromIndexes is ja        = concat [getFromIndex i ja | i <- is]

-- | Get all elements from object and array.
getChildern :: JFilter
getChildern (JObject o) = M.elems o
getChildern (JArray a)  = a
getChildern _              = []

--  Filter combinators

-- | @(f >>> g)@  - Apply filter f, later filter g .
--
-- Examples:
--
-- > query3 = getFromKeys ["Europa", "America"] >>> getFromObj 
-- > ex3 = pprints $ query3 jCitys
--
-- Result:
--
-- > [
-- >   ["Berlin", "Bonn"],["London", "Glasgow"],["Kiyv", "Gitomir", "Lviv"],["Toronto"],["NewYork"]
-- > ]
(>>>) :: JFilter -> JFilter -> JFilter 
(f >>> g) t =  concat [g t' | t' <- f t]

-- | Concat results two filters.
(<+>) :: JFilter -> JFilter -> JFilter
(f <+> g) t =  f t ++ g t

-- | @(f `orElse` g)@ - Apply f, if @f@ returned @null@ apply @g@.
orElse :: JFilter -> JFilter -> JFilter
orElse f g t
  | null res1 = g t
  | otherwise = res1
    where res1 = f t

-- | @(f `when` g)@ - When @g@ returned @not null@, apply @f@.
when :: JFilter -> JFilter -> JFilter
when f g t
  | null (g t) = [t]
  | otherwise  = f t
when' f g t = case g t of
                [] -> [t]
                _  -> f t
 
-- | @(f `guards` g )@ - If @f@ returned null then @null@ else apply @g@.
guards	:: JFilter -> JFilter -> JFilter
guards f g t
  | null (f t) = []
  | otherwise  = g t


-- | Tree traversal filter for object and array.
deep	:: JFilter -> JFilter
deep f  = f `orElse` (getChildern >>> deep f)

-- | Tree traversal filter for array.
deepObj	:: JFilter -> JFilter
deepObj f  = f `orElse` (getFromObj >>> deepObj f)

-- | Tree traversal filter for array.
--
-- Example:
--
-- >  -- Query:  All city Europa and Australia.
-- >  -- query31, query32, query33 is equal.
-- >  
-- > query31 = getFromKeys ["Europa", "Australia"] 
-- >   >>> (getFromArr `orElse`  getFromObj)
-- >   >>> (isStr `orElse` getFromArr)
-- > ex31 = pprints $ query31 jCitys
-- > 
-- > query32 = getFromKeys ["Europa", "Australia"] 
-- >   >>> (getFromObj `when` isObj)
-- >   >>> getFromArr
-- > ex32 = pprints $ query32 jCitys
-- > 
-- > query33 = getFromKeys ["Europa",  "Australia"] 
-- >   >>> deep getFromArr
-- > ex33 = pprints $ query33 jCitys
--
--   Result:
--
-- > [
-- >   "Berlin","Bonn","London","Glasgow","Kiyv","Gitomir","Lviv","Melburn","Adelaida"
-- > ]
deepArr :: JFilter -> JFilter
deepArr f  = f `orElse` (getFromArr >>> deepArr f)