-- | Utility functions for unit testing.

module Michelson.Test.Unit
  ( ContractReturn
  , ContractPropValidator
  , contractProp
  , contractPropVal
  , contractHasEntryPoints
  , matchContractEntryPoints
  , hasEp
  , validateStorageIs
  ) where

import Data.List.NonEmpty (fromList)
import qualified Data.Map as Map
import Fmt ((+|), (|+))
import Test.HUnit (Assertion, assertFailure, (@?=))

import Michelson.Interpret (ContractEnv, ContractReturn, interpret)
import Michelson.Printer (printUntypedContract)
import Michelson.Runtime (parseExpandContract)
import Michelson.Typed (Contract, IsoValue(..), ToT, epNameToParamAnn)
import qualified Michelson.Typed as T
import Michelson.Untyped (EpName, para)
import Michelson.Untyped hiding (Contract)
import qualified Michelson.Untyped as U

-- | Type for contract execution validation.
--
-- It's a function which is supplied with contract execution output
-- (failure or new storage with operation list).
--
-- Function returns a property which type is designated by type variable @prop@
-- and might be 'Test.QuickCheck.Property' or 'Test.Hspec.Expectation'
-- or anything else relevant.
type ContractPropValidator st prop = ContractReturn st -> prop

-- | Contract's property tester against given input.
-- Takes contract environment, initial storage and parameter,
-- interprets contract on this input and invokes validation function.
contractProp
  :: ( IsoValue param, IsoValue storage
     , ToT param ~ cp, ToT storage ~ st
     , T.ParameterScope cp
     )
  => Contract cp st
  -> ContractPropValidator st prop
  -> ContractEnv
  -> param
  -> storage
  -> prop
contractProp instr check env param initSt =
  contractPropVal instr check env (toVal param) (toVal initSt)

-- | Version of 'contractProp' which takes 'Val' as arguments instead
-- of regular Haskell values.
--
-- This function assumes that contract has no explicit default entrypoints
-- and you always have to construct parameter manually; if you need to test
-- contract calling specific entrypoints, use integrational testing defined
-- by "Michelson.Test.Integrational" module.
contractPropVal
  :: (T.ParameterScope cp)
  => Contract cp st
  -> ContractPropValidator st prop
  -> ContractEnv
  -> T.Value cp
  -> T.Value st
  -> prop
contractPropVal instr check env param initSt =
  check $ interpret instr T.epcCallRootUnsafe param initSt env

-- | Check if entrypoint is present in `T`.
hasEp :: T -> (EpName, U.Type) -> Bool
hasEp (TOr lFieldAnn rFieldAnn lType@(Type lT _) rType@(Type rT _))
      ep@(epNameToParamAnn -> epAnn, epType) = or
  [ (epAnn == lFieldAnn && epType == lType)
  , (epAnn == rFieldAnn && epType == rType)
  , hasEp lT ep
  , hasEp rT ep ]
hasEp _ _ = False

-- | Check whether the given set of entrypoints is present in contract.
contractHasEntryPoints :: U.Contract -> Map EpName U.Type -> Bool
contractHasEntryPoints contract eps = isRight $ matchContractEntryPoints contract eps

-- | Match the given contract with provided set of entrypoints, return left if some
-- entrypoints were not found.
matchContractEntryPoints
  :: HasCallStack
  => U.Contract
  -> Map EpName U.Type
  -> Either (NonEmpty (EpName, U.Type)) ()
matchContractEntryPoints contract eps =
  phi $ fromRight (error "Impossible") parsedContract
  where
    parsedContract = parseExpandContract Nothing
      (toText $ printUntypedContract True contract)
    phi (para -> Type t _) = conv $ filter (\ep -> not (hasEp t ep)) (Map.toList eps)
    conv l | null l = Right ()
           | otherwise = Left $ fromList l

-- | 'ContractPropValidator' that expects contract execution to
-- succeed and update storage to a particular constant value.
validateStorageIs
  :: IsoValue st
  => st -> ContractPropValidator (ToT st) Assertion
validateStorageIs expected (res, _) =
  case res of
    Left err ->
      assertFailure $ "Unexpected interpretation failure: " +| err |+ ""
    Right (_ops, got) ->
      got @?= toVal expected