{-# LANGUAGE FlexibleInstances #-}

module Data.JSON.ToGo
  ( ValueT(..), toValueT
  , matchValueT, matchValueT_
  , parseValueT, parseValueT_
  ) where

import Data.JSON.ToGo.Parser

import Data.Aeson (Value(..))
import Data.Monoid (Monoid(..))
import Data.Scientific (Scientific)
import Data.Text (Text)
import qualified Data.Vector as V
import qualified Data.HashMap.Strict as H

import Control.Applicative (Applicative, pure, (<*>))
import Control.Monad (MonadPlus, mzero, msum)
import Control.Monad.Trans.Class (lift)

data ValueT m a
  = NullM   { matchNull   ::                      m a }
  | BoolM   { matchBool   :: Bool       ->        m a }
  | NumberM { matchNumber :: Scientific ->        m a }
  | StringM { matchString :: Text       ->        m a }
  | ArrayM  { matchArray  :: Int        -> ValueT m a }
  | ObjectM { matchText   :: Text       -> ValueT m a }
  | AnyM    { matchAny    :: Value      ->        m a }
  | NoneM

toValueT :: Monad m => Value -> ValueT m Bool
toValueT Null       = NullM   $ return True
toValueT (Bool b)   = BoolM   $ return . (b ==)
toValueT (Number n) = NumberM $ return . (n ==)
toValueT (String s) = StringM $ return . (s ==)
toValueT (Array v)  = ArrayM  $ maybe NoneM toValueT . ((V.!?) v)
toValueT (Object h) = ObjectM $ maybe NoneM toValueT . (flip H.lookup h)

instance Monad m => Functor (ValueT m) where
  fmap g (NullM a)   = NullM   $         a >>= return . g
  fmap g (BoolM f)   = BoolM   $ \b -> f b >>= return . g
  fmap g (NumberM f) = NumberM $ \n -> f n >>= return . g
  fmap g (StringM f) = StringM $ \s -> f s >>= return . g
  fmap g (ArrayM f)  = ArrayM  $ fmap (fmap g) f
  fmap g (ObjectM f) = ObjectM $ fmap (fmap g) f
  fmap g (AnyM f)    = AnyM    $ \v -> f v >>= return . g
  fmap g NoneM       = NoneM

instance Monad m => Applicative (ValueT m) where
  pure = NullM . return
  (<*>) = apply

apply :: Monad m => ValueT m (a -> b) -> ValueT m a -> ValueT m b
apply (NullM g)   (NullM a)   = NullM   $ g >>= \g' -> a >>= return.g'
apply (BoolM g)   (BoolM f)   = BoolM   $ \b -> g b >>= \g' -> f b >>= return.g'
apply (NumberM g) (NumberM f) = NumberM $ \n -> g n >>= \g' -> f n >>= return.g'
apply (StringM g) (StringM f) = StringM $ \s -> g s >>= \g' -> f s >>= return.g'
apply (ArrayM g)  (ArrayM f)  = ArrayM  $ \i -> apply (g i) (f i)
apply (ObjectM g) (ObjectM f) = ObjectM $ \k -> apply (g k) (f k)
apply (AnyM g)    (AnyM f)    = AnyM    $ \v -> g v >>= \g' -> f v >>= return.g'
apply (AnyM g)    (NullM a)   = NullM   $ g Null >>= \g' -> a >>= return.g'
apply (AnyM g)    (BoolM f)   = BoolM   $ \b -> g (Bool b) >>= \g' -> f b >>= return.g'
apply (AnyM g)    (NumberM f) = NumberM $ \n -> g (Number n) >>= \g' -> f n >>= return.g'
apply (AnyM g)    (StringM f) = StringM $ \s -> g (String s) >>= \g' -> f s >>= return.g'
apply (AnyM g)    (ArrayM f)  = ArrayM  $ \i -> apply (AnyM g) (f i)
apply (AnyM g)    (ObjectM f) = ObjectM $ \k -> apply (AnyM g) (f k)
apply _           _           = NoneM

matchValueT :: MonadPlus m => ValueT m a -> Value -> m a
matchValueT (NullM ma)  Null       = ma
matchValueT (BoolM f)   (Bool b)   = f b
matchValueT (NumberM f) (Number n) = f n
matchValueT (StringM f) (String s) = f s
matchValueT (ArrayM f)  (Array v)  = msum $ map (uncurry (matchValueT . f)) (V.toList $ V.indexed v)
matchValueT (ObjectM f) (Object h) = msum $ map (uncurry (matchValueT . f)) (H.toList h)
matchValueT (AnyM f)    v          = f v
matchValueT _           _          = mzero

matchValueT_ :: Monad m => ValueT m a -> Value -> m ()
matchValueT_ (NullM ma)  Null         = ma >> return ()
matchValueT_ (BoolM f)   (Bool b)     = f b >> return ()
matchValueT_ (NumberM f) (Number n)   = f n >> return ()
matchValueT_ (StringM f) (String s)   = f s >> return ()
matchValueT_ (ArrayM f)  (Array v)    = mapM_ (uncurry (matchValueT_ . f)) (V.toList $ V.indexed v)
matchValueT_ (ObjectM f) (Object ias) = mapM_ (uncurry (matchValueT_ . f)) (H.toList ias)
matchValueT_ (AnyM f)    v            = f v >> return ()
matchValueT_ _           _            = return ()

parseValueT :: (Monad m, Monoid r) => ValueT m r -> ParserM m r
parseValueT (ArrayM f)  = parray  (parseValueT.f)
parseValueT (ObjectM f) = pobject (parseValueT.f)
parseValueT (NullM m)   = pbool   >>  lift m
parseValueT (BoolM f)   = pbool   >>= lift.f
parseValueT (NumberM f) = pnumber >>= lift.f
parseValueT (StringM f) = pstring >>= lift.f
parseValueT (AnyM f)    = pvalue  >>= lift.f
parseValueT NoneM       = fail "none"

parseValueT_ :: Monad m => ValueT m a -> ParserM m ()
parseValueT_ = parseValueT . fmap (const ())