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

module Test.Interpreter.Balance
  ( test_balanceIncludesAmount
  , test_balanceIncludesAmountComplexCase
  ) where

import Hedgehog (Gen, forAll, property, withTests)
import Hedgehog.Gen qualified as Gen
import Test.Tasty (TestTree)
import Test.Tasty.Hedgehog (testProperty)

import Morley.Michelson.Typed
import Morley.Michelson.Untyped qualified as U
import Morley.Tezos.Core
import Test.Cleveland
import Test.Cleveland.Michelson (testTreesWithUntypedContract)

import Test.Util.Contracts

data Fixture =
  Fixture
    { fStartingBalance :: Mutez
    , fAmount :: Mutez
    } deriving stock (Show)

genFixture :: Gen Fixture
genFixture = do
  fStartingBalance <- Gen.enum 1000 5000
  fAmount <- Gen.enum 0 1000
  return Fixture{..}

test_balanceIncludesAmount :: IO [TestTree]
test_balanceIncludesAmount = do
  testTreesWithUntypedContract
    (inContractsDir "check_if_balance_includes_incoming_amount.tz") $
      \checker ->
        pure
          [ testProperty "BALANCE includes AMOUNT" $ withTests 50 $ property $ do
              fixture <- forAll genFixture
              testScenarioProps $ scenario $ clevelandBalanceTestScenario
                checker fixture
          ]

clevelandBalanceTestScenario :: Monad m => U.Contract -> Fixture -> ClevelandT m ()
clevelandBalanceTestScenario checker Fixture{..} = do
  let result = unsafeAddMutez fStartingBalance fAmount

  let
    uoData = UntypedOriginateData
      { uodName = "checkIfBalanceIncludeAmount"
      , uodBalance = fStartingBalance
      , uodStorage = (untypeValue $ toVal ())
      , uodContract = checker
      }

  address <- originateUntyped uoData
  let
    transferData = TransferData
      { tdTo = address
      , tdAmount = fAmount
      , tdEntrypoint = DefEpName
      , tdParameter = result
      }

  transfer transferData
  getBalance address @@== result

test_balanceIncludesAmountComplexCase :: IO [TestTree]
test_balanceIncludesAmountComplexCase = do
  testTreesWithUntypedContract (inContractsDir "balance_test_case_a.tz") $ \contractA ->
    testTreesWithUntypedContract (inContractsDir "balance_test_case_b.tz") $ \contractB ->
      pure
        [ testScenario "BALANCE returns expected value in nested calls" $ scenario do
            let
              origDataA = UntypedOriginateData
                { uodName = "balance_test_case_a"
                , uodBalance = 0
                , uodStorage = (untypeValue $ toVal @[Mutez] [])
                , uodContract = contractA
                }
            let
              origDataB = UntypedOriginateData
                { uodName = "balance_test_case_b"
                , uodBalance = 0
                , uodStorage = (untypeValue $ toVal ())
                , uodContract = contractB
                }
            addressA <- originateUntyped origDataA
            addressB <- originateUntyped origDataB
            let
              transferData = TransferData
                { tdTo = addressA
                , tdAmount = 100
                , tdEntrypoint = DefEpName
                , tdParameter = addressB
                }

            transfer transferData

            -- A sends 30 to B, then B sends 5 back to A. A records call to BALANCE at each entry.
            -- We expect that 5 mutez sent back are included in the second call to BALANCE.
            let expectedStorage = [75, 100]
            getStorage @[Mutez] addressA @@== expectedStorage
        ]