{-|
Module      :  Data.Aeson.Schema.TH.Get
Maintainer  :  Brandon Chinn <brandon@leapyear.io>
Stability   :  experimental
Portability :  portable

The 'get' quasiquoter.
-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}

module Data.Aeson.Schema.TH.Get where

import Control.Monad ((>=>))
import Data.List (intercalate)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Maybe as Maybe
import Data.Proxy (Proxy(..))
import GHC.Stack (HasCallStack)
import Language.Haskell.TH
import Language.Haskell.TH.Quote (QuasiQuoter(..))

import Data.Aeson.Schema.Internal (getKey)
import Data.Aeson.Schema.TH.Parse
    (GetterExp(..), GetterOperation(..), GetterOps, parseGetterExp)
import Data.Aeson.Schema.Utils.Sum (fromSumType)

-- | Defines a QuasiQuoter for extracting JSON data.
--
-- Example:
--
-- > let Just result = decode ... :: Maybe (Object MySchema)
-- >
-- > [get| result.foo.a |]          :: Int
-- > [get| result.foo.nodes |]      :: [Object (..)]
-- > [get| result.foo.nodes[] |]    :: [Object (..)]
-- > [get| result.foo.nodes[].b |]  :: [Maybe Bool]
-- > [get| result.foo.nodes[].b! |] :: [Bool] -- runtime error if any values are Nothing
-- > [get| result.foo.c |]          :: Text
-- > [get| result.foo.(a,c) |]      :: (Int, Text)
-- > [get| result.foo.[c,d] |]      :: [Text]
-- >
-- > let nodes = [get| result.foo.nodes |]
-- > flip map nodes $ \node -> fromMaybe ([get| node.num |] == 0) [get| node.b |]
-- > map [get| .num |] nodes
--
-- Syntax:
--
-- * @x.y@ is only valid if @x@ is an 'Data.Aeson.Schema.Object'. Returns the value of the key @y@.
--
-- * @.y@ returns a function that takes in an 'Data.Aeson.Schema.Object' and returns the value of
--   the key @y@.
--
-- * @x.[y,z.a]@ is only valid if @x@ is an 'Data.Aeson.Schema.Object', and if @y@ and @z.a@ have
--   the same type. Returns the value of the operations @y@ and @z.a@ as a list.
--   MUST be the last operation.
--
-- * @x.(y,z.a)@ is only valid if @x@ is an 'Data.Aeson.Schema.Object'. Returns the value of the
--   operations @y@ and @z.a@ as a tuple.
--   MUST be the last operation.
--
-- * @x!@ is only valid if @x@ is a 'Maybe'. Unwraps the value of @x@ from a 'Just' value and
--   errors (at runtime!) if @x@ is 'Nothing'.
--
-- * @x[]@ is only valid if @x@ is a list. Applies the remaining rules as an 'fmap' over the
--   values in the list, e.g.
--
--     * @x[]@ without anything after is equivalent to @x@
--     * @x[].y@ gets the key @y@ in all the Objects in @x@
--     * @x[]!@ unwraps all 'Just' values in @x@ (and errors if any 'Nothing' values exist in @x@)
--
-- * @x?@ follows the same rules as @x[]@ except it's only valid if @x@ is a 'Maybe'.
--
-- * @x\@#@ is only valid if @x@ is a 'SumType'. If the sum type contains a value at the given
--   branch (e.g. @x\@0@ for @Here v@), return 'Just' that value, otherwise 'Nothing'. (added in
--   v1.1.0)
--
--   e.g. with the schema @{ a: Int | Bool }@, calling @[get| .a\@0 |]@ will return @Maybe Int@ if
--   the sum type contains an 'Int'.
get :: QuasiQuoter
get :: QuasiQuoter
get = QuasiQuoter :: (String -> Q Exp)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> Q [Dec])
-> QuasiQuoter
QuasiQuoter
  { quoteExp :: String -> Q Exp
quoteExp = String -> Q GetterExp
forall (m :: * -> *). MonadFail m => String -> m GetterExp
parseGetterExp (String -> Q GetterExp) -> (GetterExp -> Q Exp) -> String -> Q Exp
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> GetterExp -> Q Exp
generateGetterExp
  , quoteDec :: String -> Q [Dec]
quoteDec = String -> String -> Q [Dec]
forall a. HasCallStack => String -> a
error String
"Cannot use `get` for Dec"
  , quoteType :: String -> Q Type
quoteType = String -> String -> Q Type
forall a. HasCallStack => String -> a
error String
"Cannot use `get` for Type"
  , quotePat :: String -> Q Pat
quotePat = String -> String -> Q Pat
forall a. HasCallStack => String -> a
error String
"Cannot use `get` for Pat"
  }

generateGetterExp :: GetterExp -> ExpQ
generateGetterExp :: GetterExp -> Q Exp
generateGetterExp GetterExp{Maybe String
GetterOps
$sel:getterOps:GetterExp :: GetterExp -> GetterOps
$sel:start:GetterExp :: GetterExp -> Maybe String
getterOps :: GetterOps
start :: Maybe String
..} = Q Exp -> Q Exp
applyStart (Q Exp -> Q Exp) -> Q Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ GetterOpExps -> Q Exp
resolveGetterOpExps (GetterOpExps -> Q Exp) -> GetterOpExps -> Q Exp
forall a b. (a -> b) -> a -> b
$ [GetterOperation] -> GetterOps -> GetterOpExps
mkGetterOpExps [] GetterOps
getterOps
  where
    applyStart :: Q Exp -> Q Exp
applyStart Q Exp
expr = Q Exp -> (String -> Q Exp) -> Maybe String -> Q Exp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Q Exp
expr (Q Exp -> Q Exp -> Q Exp
appE Q Exp
expr (Q Exp -> Q Exp) -> (String -> Q Exp) -> String -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Q Exp
varE (Name -> Q Exp) -> (String -> Name) -> String -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Name
mkName) Maybe String
start

    startDisplay :: String
startDisplay = case Maybe String
start of
      Maybe String
Nothing -> String
""
      Just String
s -> if Char
'.' Char -> String -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
s then String
"(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")" else String
s

    mkGetterOpExps :: [GetterOperation] -> GetterOps -> GetterOpExps
    mkGetterOpExps :: [GetterOperation] -> GetterOps -> GetterOpExps
mkGetterOpExps [GetterOperation]
historyPrefix = ([GetterOperation] -> GetterOperation -> GetterOpExp)
-> GetterOps -> GetterOpExps
forall a b. ([a] -> a -> b) -> NonEmpty a -> NonEmpty b
mapWithHistory ([GetterOperation] -> GetterOperation -> GetterOpExp
mkGetterOpExp ([GetterOperation] -> GetterOperation -> GetterOpExp)
-> ([GetterOperation] -> [GetterOperation])
-> [GetterOperation]
-> GetterOperation
-> GetterOpExp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([GetterOperation]
historyPrefix [GetterOperation] -> [GetterOperation] -> [GetterOperation]
forall a. [a] -> [a] -> [a]
++))

    mkGetterOpExp :: [GetterOperation] -> GetterOperation -> GetterOpExp
    mkGetterOpExp :: [GetterOperation] -> GetterOperation -> GetterOpExp
mkGetterOpExp [GetterOperation]
history = \case
      GetterKey String
key ->
        let keyType :: Q Type
keyType = TyLitQ -> Q Type
litT (TyLitQ -> Q Type) -> TyLitQ -> Q Type
forall a b. (a -> b) -> a -> b
$ String -> TyLitQ
strTyLit String
key
        in Q Exp -> GetterOpExp
ApplyOp [| getKey (Proxy :: Proxy $keyType) |]

      GetterOperation
GetterBang ->
        let expr :: String
expr = String
startDisplay String -> String -> String
forall a. [a] -> [a] -> [a]
++ [GetterOperation] -> String
forall (t :: * -> *). Foldable t => t GetterOperation -> String
showGetterOps [GetterOperation]
history
        in Q Exp -> GetterOpExp
ApplyOp [| fromJust expr |]

      GetterOperation
GetterMapMaybe ->
        Q Exp -> GetterOpExp
ApplyOpInfix [| (<$?>) |]

      GetterOperation
GetterMapList ->
        Q Exp -> GetterOpExp
ApplyOpInfix [| (<$:>) |]

      GetterBranch Int
branch ->
        let branchType :: Q Type
branchType = TyLitQ -> Q Type
litT (TyLitQ -> Q Type) -> TyLitQ -> Q Type
forall a b. (a -> b) -> a -> b
$ Integer -> TyLitQ
numTyLit (Integer -> TyLitQ) -> Integer -> TyLitQ
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
branch
        in Q Exp -> GetterOpExp
ApplyOp [| fromSumType (Proxy :: Proxy $branchType) |]

      GetterList NonEmpty GetterOps
elemOps ->
        NonEmpty GetterOpExps -> GetterOpExp
ApplyOpsIntoList (NonEmpty GetterOpExps -> GetterOpExp)
-> NonEmpty GetterOpExps -> GetterOpExp
forall a b. (a -> b) -> a -> b
$ [GetterOperation] -> GetterOps -> GetterOpExps
mkGetterOpExps [GetterOperation]
history (GetterOps -> GetterOpExps)
-> NonEmpty GetterOps -> NonEmpty GetterOpExps
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonEmpty GetterOps
elemOps

      GetterTuple NonEmpty GetterOps
elemOps ->
        NonEmpty GetterOpExps -> GetterOpExp
ApplyOpsIntoTuple (NonEmpty GetterOpExps -> GetterOpExp)
-> NonEmpty GetterOpExps -> GetterOpExp
forall a b. (a -> b) -> a -> b
$ [GetterOperation] -> GetterOps -> GetterOpExps
mkGetterOpExps [GetterOperation]
history (GetterOps -> GetterOpExps)
-> NonEmpty GetterOps -> NonEmpty GetterOpExps
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonEmpty GetterOps
elemOps

{- Runtime helpers -}

-- | fromJust with helpful error message
fromJust :: HasCallStack => String -> Maybe a -> a
fromJust :: String -> Maybe a -> a
fromJust String
expr = a -> Maybe a -> a
forall a. a -> Maybe a -> a
Maybe.fromMaybe (String -> a
forall a. HasCallStack => String -> a
error String
errMsg)
  where
    errMsg :: String
errMsg = String
"Called 'fromJust' on null expression" String -> String -> String
forall a. [a] -> [a] -> [a]
++ if String -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
expr then String
"" else String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
expr

-- | fmap specialized to Maybe
(<$?>) :: (a -> b) -> Maybe a -> Maybe b
<$?> :: (a -> b) -> Maybe a -> Maybe b
(<$?>) = (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
(<$>)

-- | fmap specialized to [a]
(<$:>) :: (a -> b) -> [a] -> [b]
<$:> :: (a -> b) -> [a] -> [b]
(<$:>) = (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
(<$>)

{- Code generation helpers -}

data GetterOpExp
  = ApplyOp ExpQ                              -- ^ next . f
  | ApplyOpInfix ExpQ                         -- ^ (next `f`)
  | ApplyOpsIntoList (NonEmpty GetterOpExps)  -- ^ \v -> [f1 v, f2 v, ...]
  | ApplyOpsIntoTuple (NonEmpty GetterOpExps) -- ^ \v -> (f1 v, f2 v, ...)

type GetterOpExps = NonEmpty GetterOpExp

resolveGetterOpExps :: GetterOpExps -> ExpQ
resolveGetterOpExps :: GetterOpExps -> Q Exp
resolveGetterOpExps (GetterOpExp
op NonEmpty.:| [GetterOpExp]
ops) =
  case GetterOpExp
op of
    ApplyOp Q Exp
f -> [| $next . $f |]
    ApplyOpInfix Q Exp
f -> Maybe (Q Exp) -> Q Exp -> Maybe (Q Exp) -> Q Exp
infixE (Q Exp -> Maybe (Q Exp)
forall a. a -> Maybe a
Just Q Exp
next) Q Exp
f Maybe (Q Exp)
forall a. Maybe a
Nothing

    -- suffixes; ops should be empty
    ApplyOpsIntoList NonEmpty GetterOpExps
elemOps -> ([Q Exp] -> Q Exp) -> NonEmpty GetterOpExps -> Q Exp
resolveEach [Q Exp] -> Q Exp
listE NonEmpty GetterOpExps
elemOps
    ApplyOpsIntoTuple NonEmpty GetterOpExps
elemOps -> ([Q Exp] -> Q Exp) -> NonEmpty GetterOpExps -> Q Exp
resolveEach [Q Exp] -> Q Exp
tupE NonEmpty GetterOpExps
elemOps
  where
    next :: Q Exp
next = Q Exp -> (GetterOpExps -> Q Exp) -> Maybe GetterOpExps -> Q Exp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [| id |] GetterOpExps -> Q Exp
resolveGetterOpExps (Maybe GetterOpExps -> Q Exp) -> Maybe GetterOpExps -> Q Exp
forall a b. (a -> b) -> a -> b
$ [GetterOpExp] -> Maybe GetterOpExps
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [GetterOpExp]
ops

    resolveEach :: ([Q Exp] -> Q Exp) -> NonEmpty GetterOpExps -> Q Exp
resolveEach [Q Exp] -> Q Exp
fromElems NonEmpty GetterOpExps
elemOps = do
      Name
val <- String -> Q Name
newName String
"v"
      let applyVal :: Q Exp -> Q Exp
applyVal Q Exp
expr = Q Exp -> Q Exp -> Q Exp
appE Q Exp
expr (Name -> Q Exp
varE Name
val)
      [Q Pat] -> Q Exp -> Q Exp
lamE [Name -> Q Pat
varP Name
val] (Q Exp -> Q Exp) -> Q Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Q Exp] -> Q Exp
fromElems ([Q Exp] -> Q Exp) -> [Q Exp] -> Q Exp
forall a b. (a -> b) -> a -> b
$ (GetterOpExps -> Q Exp) -> [GetterOpExps] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Q Exp -> Q Exp
applyVal (Q Exp -> Q Exp)
-> (GetterOpExps -> Q Exp) -> GetterOpExps -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GetterOpExps -> Q Exp
resolveGetterOpExps) ([GetterOpExps] -> [Q Exp]) -> [GetterOpExps] -> [Q Exp]
forall a b. (a -> b) -> a -> b
$ NonEmpty GetterOpExps -> [GetterOpExps]
forall a. NonEmpty a -> [a]
NonEmpty.toList NonEmpty GetterOpExps
elemOps

showGetterOps :: Foldable t => t GetterOperation -> String
showGetterOps :: t GetterOperation -> String
showGetterOps = (GetterOperation -> String) -> t GetterOperation -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap GetterOperation -> String
showGetterOp
  where
    showGetterOp :: GetterOperation -> String
showGetterOp = \case
      GetterKey String
key -> Char
'.'Char -> String -> String
forall a. a -> [a] -> [a]
:String
key
      GetterOperation
GetterBang -> String
"!"
      GetterOperation
GetterMapList -> String
"[]"
      GetterOperation
GetterMapMaybe -> String
"?"
      GetterBranch Int
x -> Char
'@' Char -> String -> String
forall a. a -> [a] -> [a]
: Int -> String
forall a. Show a => a -> String
show Int
x
      GetterList NonEmpty GetterOps
elemOps -> String
".[" String -> String -> String
forall a. [a] -> [a] -> [a]
++ NonEmpty GetterOps -> String
showGetterOpsList NonEmpty GetterOps
elemOps String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"]"
      GetterTuple NonEmpty GetterOps
elemOps -> String
".(" String -> String -> String
forall a. [a] -> [a] -> [a]
++ NonEmpty GetterOps -> String
showGetterOpsList NonEmpty GetterOps
elemOps String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
")"

    showGetterOpsList :: NonEmpty GetterOps -> String
showGetterOpsList = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"," ([String] -> String)
-> (NonEmpty GetterOps -> [String]) -> NonEmpty GetterOps -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty String -> [String]
forall a. NonEmpty a -> [a]
NonEmpty.toList (NonEmpty String -> [String])
-> (NonEmpty GetterOps -> NonEmpty String)
-> NonEmpty GetterOps
-> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (GetterOps -> String) -> NonEmpty GetterOps -> NonEmpty String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap GetterOps -> String
forall (t :: * -> *). Foldable t => t GetterOperation -> String
showGetterOps

{- Utilities -}

-- | Run the given function for each element in the list, providing all elements seen so far.
--
-- e.g. for a list [1,2,3], this will return the result of
--
--   [f [] 1, f [1] 2, f [1,2] 3]
mapWithHistory :: ([a] -> a -> b) -> NonEmpty a -> NonEmpty b
mapWithHistory :: ([a] -> a -> b) -> NonEmpty a -> NonEmpty b
mapWithHistory [a] -> a -> b
f NonEmpty a
xs = ([a] -> a -> b) -> NonEmpty [a] -> NonEmpty a -> NonEmpty b
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
NonEmpty.zipWith [a] -> a -> b
f (NonEmpty a -> NonEmpty [a]
forall (f :: * -> *) a. Foldable f => f a -> NonEmpty [a]
NonEmpty.inits NonEmpty a
xs) NonEmpty a
xs