{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}

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

The 'get' quasiquoter.
-}
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
    { quoteExp :: [Char] -> Q Exp
quoteExp = forall (m :: * -> *). MonadFail m => [Char] -> m GetterExp
parseGetterExp forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> GetterExp -> Q Exp
generateGetterExp
    , quoteDec :: [Char] -> Q [Dec]
quoteDec = forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot use `get` for Dec"
    , quoteType :: [Char] -> Q Type
quoteType = forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot use `get` for Type"
    , quotePat :: [Char] -> Q Pat
quotePat = forall a. HasCallStack => [Char] -> a
error [Char]
"Cannot use `get` for Pat"
    }

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

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

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

    mkGetterOpExp :: [GetterOperation] -> GetterOperation -> GetterOpExp
    mkGetterOpExp :: [GetterOperation] -> GetterOperation -> GetterOpExp
mkGetterOpExp [GetterOperation]
history = \case
      GetterKey [Char]
key ->
        let keyType :: Q Type
keyType = forall (m :: * -> *). Quote m => m TyLit -> m Type
litT forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => [Char] -> m TyLit
strTyLit [Char]
key
         in Q Exp -> GetterOpExp
ApplyOp [|getKey (Proxy :: Proxy $keyType)|]
      GetterOperation
GetterBang ->
        let expr :: [Char]
expr = [Char]
startDisplay forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *). Foldable t => t GetterOperation -> [Char]
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 = forall (m :: * -> *). Quote m => m TyLit -> m Type
litT forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => Integer -> m TyLit
numTyLit forall a b. (a -> b) -> a -> b
$ 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 forall a b. (a -> b) -> a -> b
$ [GetterOperation] -> GetterOps -> GetterOpExps
mkGetterOpExps [GetterOperation]
history forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NonEmpty GetterOps
elemOps
      GetterTuple NonEmpty GetterOps
elemOps ->
        NonEmpty GetterOpExps -> GetterOpExp
ApplyOpsIntoTuple forall a b. (a -> b) -> a -> b
$ [GetterOperation] -> GetterOps -> GetterOpExps
mkGetterOpExps [GetterOperation]
history 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 :: forall a. HasCallStack => [Char] -> Maybe a -> a
fromJust [Char]
expr = forall a. a -> Maybe a -> a
Maybe.fromMaybe (forall a. HasCallStack => [Char] -> a
error [Char]
errMsg)
  where
    errMsg :: [Char]
errMsg = [Char]
"Called 'fromJust' on null expression" forall a. [a] -> [a] -> [a]
++ if forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
expr then [Char]
"" else [Char]
": " forall a. [a] -> [a] -> [a]
++ [Char]
expr

-- | fmap specialized to Maybe
(<$?>) :: (a -> b) -> Maybe a -> Maybe b
<$?> :: forall a 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]
<$:> :: forall a b. (a -> b) -> [a] -> [b]
(<$:>) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
(<$>)

{- Code generation helpers -}

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

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 -> forall (m :: * -> *).
Quote m =>
Maybe (m Exp) -> m Exp -> Maybe (m Exp) -> m Exp
infixE (forall a. a -> Maybe a
Just Q Exp
next) Q Exp
f forall a. Maybe a
Nothing
    -- suffixes; ops should be empty
    ApplyOpsIntoList NonEmpty GetterOpExps
elemOps -> forall {m :: * -> *}.
Quote m =>
([Q Exp] -> m Exp) -> NonEmpty GetterOpExps -> m Exp
resolveEach forall (m :: * -> *). Quote m => [m Exp] -> m Exp
listE NonEmpty GetterOpExps
elemOps
    ApplyOpsIntoTuple NonEmpty GetterOpExps
elemOps -> forall {m :: * -> *}.
Quote m =>
([Q Exp] -> m Exp) -> NonEmpty GetterOpExps -> m Exp
resolveEach forall (m :: * -> *). Quote m => [m Exp] -> m Exp
tupE NonEmpty GetterOpExps
elemOps
  where
    next :: Q Exp
next = forall b a. b -> (a -> b) -> Maybe a -> b
maybe [|id|] GetterOpExps -> Q Exp
resolveGetterOpExps forall a b. (a -> b) -> a -> b
$ forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [GetterOpExp]
ops

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

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

    showGetterOpsList :: NonEmpty GetterOps -> [Char]
showGetterOpsList = forall a. [a] -> [[a]] -> [a]
intercalate [Char]
"," forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NonEmpty a -> [a]
NonEmpty.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *). Foldable t => t GetterOperation -> [Char]
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 :: forall a b. ([a] -> a -> b) -> NonEmpty a -> NonEmpty b
mapWithHistory [a] -> a -> b
f NonEmpty a
xs = forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
NonEmpty.zipWith [a] -> a -> b
f (forall (f :: * -> *) a. Foldable f => f a -> NonEmpty [a]
NonEmpty.inits NonEmpty a
xs) NonEmpty a
xs