{-
  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 TestCommon where

import Control.Exception (throw, throwIO, evaluate)
import Control.Monad
import Control.Monad.Trans.Class
import Data.IORef
import Test.HUnit hiding (State)

import Thrift.Api
import Thrift.Channel
import Thrift.Monad
import Thrift.Protocol.ApplicationException.Types

import Math.Types
import Math.Adder.Client
import Math.Adder.Service
import Math.Calculator.Client
import Math.Calculator.Service

-- Server Implementation for the 'Calculator' service

type State = IORef Int

initServerState :: IO State
initServerState = newIORef 0

processCommand :: State -> CalculatorCommand a -> IO a
processCommand _ (SuperAdder (Add x y)) = pure $ x + y
processCommand _ (Divide x y)
  | y == 0    = throwIO DivideByZero
  | otherwise = pure $ x / y
processCommand _ (QuotRem x y) = pure $ QuotRemResponse (x `quot` y) (x `rem` y)
processCommand state (Put x) = writeIORef state x
processCommand state (PutMany xs) = mapM_ (writeIORef state) xs
processCommand state Get = readIORef state
processCommand _ Unimplemented =
  throw $ ApplicationException "" ApplicationExceptionType_UnknownMethod


-- common client computations/tests used with diffeerent channel implementations

-- if needed, this 'Calculator' could just become a type
-- parameter, in case such utilities are needed for other
-- test suites.
data ChannelTest = ChannelTest
  { ctestName :: String
  , ctestOpts :: RpcOptions
  , ctestAct  :: Thrift Calculator ()
  }

-- | Package up a test label and the corresponding client
--   computation, to later be executed against a specific
--   channel implementation and server.
channelTest
  :: String -> Thrift Calculator () -> ChannelTest
channelTest lbl comp = channelTestWithOpts lbl defaultRpcOptions comp

-- | Like 'channelTest', but with the ability to pass 'RpcOptions' to the
--   test.
channelTestWithOpts
  :: String -> RpcOptions -> Thrift Calculator () -> ChannelTest
channelTestWithOpts = ChannelTest

runChannelTests
  :: (String -> RpcOptions -> Thrift Calculator () -> Test)
  -> [ChannelTest]
  -> [Test]
runChannelTests testFun tests =
  map (\ChannelTest{..} -> testFun ctestName ctestOpts ctestAct) tests

addTest :: ChannelTest
addTest = channelTest "add test" $ do
  result <- add 5 2
  lift $ assertEqual "5 + 2 = 7" 7 result

divTest :: ChannelTest
divTest = channelTest "divide test" $ do
  result <- divide 10 2
  lift $ assertEqual "10 / 2 = 5" 5 result

putGetTest :: ChannelTest
putGetTest = channelTest "put get" $ do
  let value = 99
  put value
  result <- get
  lift $ assertEqual "put 99 = get" value result

putPutGetTest :: ChannelTest
putPutGetTest = channelTest "put put get" $ do
  let value1 = 99
      value2 = 100
  put value1
  put value2
  result <- get
  lift $ assertEqual "put 99 then put 100 = get (= put 100)" value2 result

exceptionTest :: ChannelTest
exceptionTest = channelTest "exception test" $
  (void . lift . evaluate =<< divide 0 0)
    `catchThrift` \DivideByZero -> return ()

optionsTest :: ChannelTest
optionsTest = channelTestWithOpts "options test" opts $
  (void . lift . evaluate =<< (add 0 0 *> error "fail"))
    `catchThrift` \ChannelException{} -> return ()
  where
    opts = setRpcPriority defaultRpcOptions High

unimplementedTest :: ChannelTest
unimplementedTest = channelTest "unimplemented test" $
  unimplemented `catchThrift` \ApplicationException{} -> return ()

multiTest :: ChannelTest
multiTest = channelTest "multiple requests" $ do
  r1 <- add 2 2
  lift $ assertEqual "2 + 2 = 4" 4 r1

  r2 <- divide 64 16
  lift $ assertEqual "64 / 16 = 4" 4 r2

  put 100
  r3 <- get
  lift $ assertEqual "put = get" 100 r3