content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# LANGUAGE NoImplicitPrelude #-}
  
  module Data.Morpheus.App.Internal.Resolving.Batching
    ( ResolverMapT (..),
      SelectionRef,
      runBatchedT,
      MonadBatching (..),
    )
  where
  
  import Control.Monad.Except (MonadError (throwError))
  import Data.ByteString.Lazy.Char8 (unpack)
  import Data.HashMap.Lazy (keys)
  import Data.Morpheus.App.Internal.Resolving.Cache
    ( CacheKey (..),
      CacheT,
      CacheValue (..),
      cacheResolverValues,
      cacheValue,
      isCached,
      printSelectionKey,
      useCached,
      withDebug,
    )
  import Data.Morpheus.App.Internal.Resolving.Refs (scanRefs)
  import Data.Morpheus.App.Internal.Resolving.ResolverState (ResolverContext)
  import Data.Morpheus.App.Internal.Resolving.Types
    ( NamedResolver (..),
      NamedResolverResult (..),
      ResolverMap,
    )
  import Data.Morpheus.App.Internal.Resolving.Utils
    ( NamedResolverRef (..),
      ResolverMonad,
      ResolverValue (ResEnum, ResNull, ResObject, ResRef, ResScalar),
    )
  import Data.Morpheus.Core (render)
  import Data.Morpheus.Internal.Utils (Empty (empty), IsMap (..), selectOr)
  import Data.Morpheus.Types.Internal.AST
    ( GQLError,
      Msg (..),
      SelectionContent,
      TypeName,
      VALID,
      ValidValue,
      internal,
    )
  import GHC.Show (Show (show))
  import Relude hiding (empty, show)
  
  data BatchEntry = BatchEntry
    { batchedSelection :: SelectionContent VALID,
      batchedType :: TypeName,
      batchedArguments :: [ValidValue]
    }
  
  instance Show BatchEntry where
    show BatchEntry {..} =
      "\nBATCH("
        <> toString batchedType
        <> "):"
        <> "\n  sel:"
        <> printSelectionKey batchedSelection
        <> "\n  dep:"
        <> show (map (unpack . render) batchedArguments)
  
  type SelectionRef = (SelectionContent VALID, NamedResolverRef)
  
  uniq :: (Eq a, Hashable a) => [a] -> [a]
  uniq = keys . unsafeFromList . map (,True)
  
  buildBatches :: [SelectionRef] -> [BatchEntry]
  buildBatches inputs =
    let entityTypes = uniq $ map (second resolverTypeName) inputs
     in mapMaybe (selectByEntity inputs) entityTypes
  
  selectByEntity :: [SelectionRef] -> (SelectionContent VALID, TypeName) -> Maybe BatchEntry
  selectByEntity inputs (tSel, tName) = case gerArgs (filter areEq inputs) of
    [] -> Nothing
    args -> Just (BatchEntry tSel tName args)
      where
  
    where
      gerArgs = uniq . concatMap (resolverArgument . snd)
      areEq (sel, v) = sel == tSel && tName == resolverTypeName v
  
  newtype ResolverMapT m a = ResolverMapT
    { _runResMapT :: ReaderT (ResolverMap m) (CacheT m) a
    }
    deriving
      ( Functor,
        Applicative,
        Monad
      )
  
  instance (MonadReader ResolverContext m) => MonadReader ResolverContext (ResolverMapT m) where
    ask = ResolverMapT (lift ask)
    local f (ResolverMapT m) = ResolverMapT (ReaderT (local f . runReaderT m))
  
  instance MonadTrans ResolverMapT where
    lift = ResolverMapT . lift . lift
  
  deriving instance (MonadError GQLError m) => MonadError GQLError (ResolverMapT m)
  
  runBatchedT :: (Monad m) => ResolverMapT m a -> ResolverMap m -> m a
  runBatchedT (ResolverMapT m) rmap = fst <$> runStateT (runReaderT m rmap) empty
  
  toKeys :: BatchEntry -> [CacheKey]
  toKeys (BatchEntry sel name deps) = map (CacheKey sel name) deps
  
  inCache :: (Monad m) => CacheT m a -> ResolverMapT m a
  inCache = ResolverMapT . lift
  
  class (MonadTrans t) => MonadBatching t where
    resolveRef :: (ResolverMonad m) => SelectionContent VALID -> NamedResolverRef -> t m (CacheKey, CacheValue m)
    storeValue :: (ResolverMonad m) => CacheKey -> ValidValue -> t m ValidValue
  
  instance MonadBatching IdentityT where
    resolveRef _ _ = throwError $ internal "batching is only allowed with named resolvers"
    storeValue _ _ = throwError $ internal "batching is only allowed with named resolvers"
  
  instance MonadBatching ResolverMapT where
    resolveRef sel (NamedResolverRef typename [arg]) = do
      let key = CacheKey sel typename arg
      alreadyCached <- inCache (isCached key)
      if alreadyCached
        then pure ()
        else prefetch (BatchEntry sel typename [arg])
      inCache $ do
        value <- useCached key
        pure (key, value)
    resolveRef _ ref = throwError (internal ("expected only one resolved value for " <> msg (show ref :: String)))
    storeValue key = inCache . cacheValue key
  
  prefetch :: (ResolverMonad m) => BatchEntry -> ResolverMapT m ()
  prefetch batch = do
    value <- run batch
    batches <- buildBatches . concat <$> traverse (lift . scanRefs (batchedSelection batch)) value
    resolvedEntries <- traverse (\b -> (b,) <$> run b) batches
    let caches = foldMap zipCaches $ (batch, value) : resolvedEntries
    inCache $ cacheResolverValues caches
    where
      zipCaches (b, res) = zip (toKeys b) res
      run = withDebug >=> runBatch
  
  runBatch :: (MonadError GQLError m, MonadReader ResolverContext m) => BatchEntry -> ResolverMapT m [ResolverValue m]
  runBatch (BatchEntry _ name deps)
    | null deps = pure []
    | otherwise = do
        resolvers <- ResolverMapT ask
        NamedResolver {resolverFun} <- lift (selectOr notFound pure name resolvers)
        map (toResolverValue name) <$> lift (resolverFun deps)
    where
      notFound = throwError ("resolver type " <> msg name <> "can't found")
  
  toResolverValue :: (Monad m) => TypeName -> NamedResolverResult m -> ResolverValue m
  toResolverValue typeName (NamedObjectResolver res) = ResObject (Just typeName) res
  toResolverValue _ (NamedUnionResolver unionRef) = ResRef $ pure unionRef
  toResolverValue _ (NamedEnumResolver value) = ResEnum value
  toResolverValue _ NamedNullResolver = ResNull
  toResolverValue _ (NamedScalarResolver v) = ResScalar v
  