{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

module Glean.Test.Mock
  ( IOBased(..)
  , Mock
  , mock
  , implement
  , call
  , augment
  , reimplement
  , prepare
  , doBefore
  , doFinally
  , unexpectedCall
  ) where

import Control.Exception (finally)
import Control.Monad (join)
import Data.IORef

-- | Things of type a1 -> ... -> an -> IO b
class IOBased a where
  -- | Constant action
  constIO :: (forall b. IO b) -> a

  -- | Transform the `IO` action
  mapIO :: (forall b. IO b -> IO b) -> a -> a

  -- | Push an outer `IO` into the `IO` action
  joinIO :: IO a -> a

instance IOBased (IO a) where
  constIO io = io
  mapIO f = f
  joinIO = join

instance IOBased b => IOBased (a -> b) where
  constIO io _ = constIO io
  mapIO f g x = mapIO f (g x)
  joinIO io x = joinIO $ ($ x) <$> io

-- | Execute an `IO` before the action
doBefore :: IOBased a => IO () -> a -> a
doBefore p = mapIO (p >>)

-- | Execute an `IO` after the action
doFinally :: IOBased a => IO () -> a -> a
doFinally p = mapIO (`finally` p)

infixr 5 :>
data Stream a = a :> Stream a

instance Functor Stream where
  fmap f (x :> xs) = f x :> fmap f xs

instance Applicative Stream where
  pure x = let xs = x :> xs in xs
  (f :> fs) <*> (x :> xs) = f x :> (fs <*> xs)

prepend :: [a] -> Stream a -> Stream a
prepend xs s = foldr (:>) s xs

-- | A mocked action.
newtype Mock a = Mock (IORef (Stream a))

-- | Produce an "unexpected call" error
unexpectedCall :: IOBased a => String -> a
unexpectedCall name = constIO $ fail $ "unexpected call to '" ++ name ++ "'"

-- | Mock an `IOBased` action. All calls will produce "unexpected call" errors.
mock :: IOBased a => String -> IO (Mock a)
mock name = implement name $ unexpectedCall name

-- | Produce a 'Mock' based on an implementation of the operation. The name is
-- currently ignored but will be useful later.
implement :: String -> a -> IO (Mock a)
implement _name x = Mock <$> newIORef (pure x)

-- | Call a mocked 'IOBased' action.
call :: IOBased a => Mock a -> a
call (Mock ref) = joinIO $ atomicModifyIORef ref $ \(x :> xs) -> (xs, x)

augment :: Mock a -> (a -> a) -> IO ()
augment (Mock ref) f = atomicModifyIORef ref $ \xs -> (f <$> xs, ())

reimplement :: Mock a -> a -> IO ()
reimplement m = augment m . const

-- | Transform the next calls to the 'Mock'.
prepare :: Mock a -> [a -> a] -> IO ()
prepare (Mock ref) fs = atomicModifyIORef ref $ \xs ->
  (prepend fs (pure id) <*> xs, ())