-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

module Morley.Util.Batching
  ( BatchingM
  , runBatching
  , unsafeRunBatching
  , submitThenParse

  , BatchingError (..)
  ) where

import Control.Monad.Except (Except, runExcept, throwError)
import Fmt (Buildable(..))

-- | Errors that can occur during batching, usually because the
-- underlying function that performs batch operation returns output
-- that does not match the provided input.
data BatchingError e
  -- | The function that executes the batch returned less elements in
  -- output than were provided at input.
  = InsufficientOutput
  -- | The function that executes the batch returned more elements in
  -- output than were provided at input.
  | ExtraOutput
  -- | User-provided parsing method failed.
  -- Usually this means that output does not correspond to provided input.
  | UnexpectedElement e

instance Buildable e => Buildable (BatchingError e) where
  build = \case
    InsufficientOutput ->
      "Too few elements in output of batch operation"
    ExtraOutput ->
      "Too many elements in output of batch operation"
    UnexpectedElement e ->
      "Unexpected element: " <> build e

-- | Records operations to be executed in batch.
--
-- Chronologically, this works in 3 steps:
--
-- * Form the list of input items @i@;
-- * Perform the batch operation;
-- * Parse output items @o@ into result @a@, maybe producing error @e@.
--
-- However in code we usually want steps 1 and 3 to be grouped
-- and step 2 to be delayed - 'BatchingM' facilitates this separation.
--
-- Note that 'BatchingM' is fundamentally not a monad, rather just an applicative,
-- because within a batch you cannot use result of one operation in another
-- operation.
data BatchingM i o e a = BatchingM
  { bInput :: Endo [i]
    -- ^ All the provided input, in some sort of DList
  , bParseOutput :: StateT [o] (Except (BatchingError e)) a
    -- ^ Parser for output when it is available
  } deriving stock Functor

instance Applicative (BatchingM i o e) where
  pure a = BatchingM
    { bInput = mempty, bParseOutput = pure a }
  b1 <*> b2 = BatchingM
    { bInput = bInput b1 <> bInput b2
    , bParseOutput = bParseOutput b1 <*> bParseOutput b2
    }

-- | Run recorded operations sequence using the given batch executor.
runBatching
  :: (Functor m)
  => ([i] -> m (r, [o]))
  -> BatchingM i o e a
  -> m (r, Either (BatchingError e) a)
runBatching execBatch BatchingM{..} =
    second parseResult <$> execBatch (appEndo bInput [])
  where
    parseResult output =
      runExcept (runStateT bParseOutput output) >>= \case
        (a, []) -> pure a
        _ -> throwError ExtraOutput

-- | Similar to 'runBatching', for cases when the given batch executor
-- is guaranteed to return the output respective to the provided input.
unsafeRunBatching
  :: (Functor m, Buildable e)
  => ([i] -> m (r, [o]))
  -> BatchingM i o e a
  -> m (r, a)
unsafeRunBatching =
  fmap (second unsafe) ... runBatching

-- | This is the basic primitive for all actions in 'BatchingM'.
--
-- It records that given input item should be put to batch, and once operation
-- is actually performed, the result should be parsed with given method.
submitThenParse :: i -> (o -> Either e a) -> BatchingM i o e a
submitThenParse inp parse = BatchingM
  { bInput = Endo (inp :)
  , bParseOutput = StateT $ \case
      [] -> throwError InsufficientOutput
      (o : os) -> case parse o of
        Left e -> throwError $ UnexpectedElement e
        Right x -> pure (x, os)
  }
infix 1 `submitThenParse`