-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Utils for testing Indigo

module Test.Util
  ( testIndigoContract
  , testIndigo
  , testIndigoDoc

  , zeroDivFail
  , notNewKeyFail
  , notNewKeyM
  , negativeResFail
  , negativeResM

  , validateContract
  , validateContractOps
  , validateContractSt
  , validateContractConst

  , validateStSuccess
  , validateStEither
  , validateStack2

  , noOptimizationContract
  ) where

import qualified Data.Text.IO.Utf8 as Utf8 (readFile)
import Fmt (pretty)
import Hedgehog (Gen, MonadTest, PropertyT, annotate, forAll, property, (===))
import Prelude
import Test.HUnit ((@?=))
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Hedgehog (testProperty)
import Test.Tasty.HUnit (testCase)

import Indigo.Lorentz
import Lorentz.Run (Contract(..))
import Lorentz.Test (ContractPropValidator, contractProp, dummyContractEnv, failedTest)
import Michelson.Interpret (MichelsonFailed(..))
import Michelson.Typed.Haskell.Value (IsoValuesStack)

type IndigoInstrValidator m pm st out =
  pm -> st -> Either MichelsonFailed (Rec Identity out) -> m ()

--------------------------------------------------------------------------------
-- Tests
--------------------------------------------------------------------------------

-- | Takes a validating function, an Indigo-generated contract and a Lorentz
-- contract and checks that the two have equivalent code (aka made of the same
-- instructions) and that it passes validation for random param and storage.
testIndigoContract
  :: forall pm st.
     ( Show pm, Show st
     , NiceParameterFull pm, NiceStorage st
     )
  => String
  -> Gen pm
  -> Gen st
  -> (pm -> st -> ContractPropValidator (ToT st) (PropertyT IO ()))
  -> ContractCode pm st
  -> FilePath
  -> TestTree
testIndigoContract name genPm genSt propValidator iContract michelsonFile =
  testGroup ("Indigo contract: " <> name)
    [ testCase "matches Michelson reference contract" $ do
        expectedContract <- Utf8.readFile michelsonFile
        printLorentzContract False iContractWithoutOptimization @?= fromStrict expectedContract
    , testProperty "has the correct resulting state and operations" $ property $ do
        pm <- forAll genPm
        st <- forAll genSt
        contProp pm st
    ]
  where
    iContractWithoutOptimization = noOptimizationContract iContract

    contProp :: pm -> st -> PropertyT IO ()
    contProp param storage =
      withDict (niceParameterEvi @pm) $
      contractProp
        (compileLorentzContract iContractWithoutOptimization)
        (propValidator param storage)
        dummyContractEnv param storage

testIndigoDoc
    :: forall pm st.
       String
    -> ContractCode pm st
    -> ContractCode pm st
    -> TestTree
testIndigoDoc name iContract lContract =
  testCase (name <> " matches Lorentz docs content") $
    renderLorentzDoc iContract @?= renderLorentzDoc lContract

-- | Takes a validating function and an Indigo-generated Lorentz `Instr` to check
-- the resulting stack content.
testIndigo
  :: (Show pm, Show st, IsoValue pm, IsoValue st, IsoValuesStack out)
  => String
  -> Gen pm
  -> Gen st
  -> IndigoInstrValidator (PropertyT IO) pm st out
  -> ('[pm, st] :-> out)
  -> TestTree
testIndigo name genPm genSt validator iInstr =
  testProperty (name <> " Indigo Expr has correct resulting stack") $ property $ do
    pm <- forAll genPm
    st <- forAll genSt
    stackProp pm st
  where
    stackProp param storage = validator param storage .
      interpretLorentzInstr dummyContractEnv iInstr $
        Identity param :& Identity storage :& RNil

--------------------------------------------------------------------------------
-- Common failures
--------------------------------------------------------------------------------

zeroDivFail :: MichelsonFailed
zeroDivFail = errorToVal [mt|devision by zero|] MichelsonFailedWith

notNewKeyFail :: MichelsonFailed
notNewKeyFail = errorToVal notNewKeyM MichelsonFailedWith

notNewKeyM :: MText
notNewKeyM = [mt|not new key|]

negativeResFail :: MichelsonFailed
negativeResFail = errorToVal negativeResM MichelsonFailedWith

negativeResM :: MText
negativeResM = [mt|unacceptable negative result|]

--------------------------------------------------------------------------------
-- Contract Validators
--------------------------------------------------------------------------------

-- | Makes a validator for `testIndigoContract` that can expect a failure or a
-- resulting [Operation] from the given function. Ignores new storage value.
validateContract
  :: MonadTest m
  => IsoValue st
  => (pm -> st -> Either MichelsonFailed ([Operation], st))
  -> pm -> st -> ContractPropValidator (ToT st) (m ())
validateContract fn param st (res, _) =
  assertMichelsonResult (fn param st) res $ \(ops, val) (opsRes, resVal) -> do
    annotate "matches resulting Storage and Operations"
    (ops, toVal val) === (opsRes, resVal)


-- | Makes a validator for `testIndigoContract` that can expect a failure or a
-- resulting [Operation] from the given function. Ignores new storage value.
validateContractOps
  :: MonadTest m
  => (pm -> st -> Either MichelsonFailed [Operation])
  -> pm -> st -> ContractPropValidator (ToT st) (m ())
validateContractOps fn param st (res, _) =
  assertMichelsonResult (fn param st) res $ \ops (opsRes, _) -> do
    annotate "matches resulting Operations"
    ops === opsRes

-- | Makes a validator for `testIndigoContract` that can expect a failure or a
-- new storage from the given function. Ignores resulting [Operation]
validateContractSt
  :: MonadTest m
  => IsoValue st
  => (pm -> st -> Either MichelsonFailed st)
  -> pm -> st -> ContractPropValidator (ToT st) (m ())
validateContractSt fn param st (res, _) =
  assertMichelsonResult (fn param st) res $ \val (_, resVal) -> do
    annotate "matches resulting Storage"
    toVal val === resVal

-- | Validator for `testIndigoContract` that expects the storage to remain
-- the same and the resulting [Operation] to be empty. Ignores the parameter.
validateContractConst
  :: MonadTest m
  => IsoValue st
  => pm -> st -> ContractPropValidator (ToT st) (m ())
validateContractConst = validateContract (\_param st -> Right ([], st))

--------------------------------------------------------------------------------
-- Instr Validators
--------------------------------------------------------------------------------

-- | Makes a validator for `testIndigo` that expects the stack not to change type
-- and the "storage" value to have changed as described by the given function.
-- Resulting "param" is ignored.
validateStSuccess
  :: (MonadTest m, Eq st, Show st)
  => (pm -> st -> st)
  -> IndigoInstrValidator m pm st '[pm, st]
validateStSuccess fn = validateStEither (\p s -> Right $ fn p s)

-- | Makes a validator for `testIndigo` that expects the stack not to change type
-- and Either end with a failure or with a new stack. Resulting "param" is ignored.
validateStEither
  :: forall m st pm. (MonadTest m, Eq st, Show st)
  => (pm -> st -> Either MichelsonFailed st)
  -> IndigoInstrValidator m pm st '[pm, st]
validateStEither fn param st res = assertMichelsonResult (fn param st) res checkSt
  where
    checkSt :: st -> (Rec Identity '[pm, st]) -> m ()
    checkSt val resStack = do
      let Identity _ :& Identity newState :& RNil = resStack
      annotate "matches resulting state"
      val === newState

-- | Makes a validator for `testIndigo` that expects the stack to have 3 element
-- (in order and with given values) or a failure to occur.
validateStack2
  :: forall m st pm . (MonadTest m, Eq pm, Eq st, Show pm, Show st)
  => (pm -> st -> Either MichelsonFailed (pm, st))
  -> IndigoInstrValidator m pm st '[pm, st]
validateStack2 fn param st res = assertMichelsonResult (fn param st) res checkSt
  where
    checkSt :: (pm, st) -> (Rec Identity '[pm, st]) -> m ()
    checkSt val resStack = do
      let Identity newParam :& Identity newState :& RNil = resStack
      annotate "matches resulting state"
      val === (newParam, newState)


--------------------------------------------------------------------------------
-- Helpers
--------------------------------------------------------------------------------

assertMichelsonResult
  :: MonadTest m
  => Either MichelsonFailed a
  -> Either MichelsonFailed b
  -> (a -> b -> m ())
  -> m ()
assertMichelsonResult mRes1 mRes2 validatorRight = case (mRes1, mRes2) of
  (Left err, Left e) -> annotate "expected failure" >> err === e
  (Left err, Right _) -> failedTest $ "should have failed with: " <> pretty err
  (Right _, Left e) -> failedTest $ "unexpected failure: " <> pretty e
  (Right val1, Right val2) -> validatorRight val1 val2

noOptimizationContract :: ContractCode param st -> Contract param st
noOptimizationContract code = Contract
  { cCode = code
  , cDisableInitialCast = False
  , cCompilationOptions = noOptimizationOptions
  }

noOptimizationOptions :: CompilationOptions
noOptimizationOptions = defaultCompilationOptions { coOptimizerConf = Nothing }