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

import Test.QuickCheck (Arbitrary(..), choose, withMaxSuccess)
import Test.QuickCheck.Instances.Text ()
import Test.Tasty (TestTree)
import Test.Tasty.HUnit (testCase)
import Test.Tasty.QuickCheck (testProperty)

import Michelson.Runtime.GState
import Michelson.Test (testTreesWithUntypedContract)
import Michelson.Test.Integrational
import Michelson.Typed
import qualified Michelson.Untyped as U
import Tezos.Core

import Test.Util.Contracts

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

instance Arbitrary Fixture where
  arbitrary = do
    fStartingBalance <- unsafeMkMutez <$> choose (1000, 5000)
    fAmount <- unsafeMkMutez <$> choose (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" $ withMaxSuccess 50 $
              integrationalTestProperty . scenario checker
          ]
  where
    scenario :: U.Contract -> Fixture -> IntegrationalScenario
    scenario checker Fixture{..} = do
      let result = unsafeAddMutez fStartingBalance fAmount

      address <- originate
        checker
        "checkIfBalanceIncludeAmount"
        (untypeValue $ toVal ())
        fStartingBalance

      let
        txData = TxData
          { tdSenderAddress = genesisAddress
          , tdParameter = untypeValue $ toVal result
          , tdEntrypoint = DefEpName
          , tdAmount = fAmount
          }

      transfer txData address
      validate
        $ Right
        $ expectBalance 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
          [ testCase "BALANCE returns expected value in nested calls"
              $ integrationalTestExpectation
              $ scenario contractA contractB
          ]
  where
    scenario :: U.Contract -> U.Contract -> IntegrationalScenario
    scenario contractA contractB = do
      addressA <- originate
        contractA
        "balance_test_case_a"
        (untypeValue $ toVal @[Mutez] [])
        (unsafeMkMutez 0)

      addressB <- originate
        contractB
        "balance_test_case_b"
        (untypeValue $ toVal ())
        (unsafeMkMutez 0)

      let
        txData = TxData
          { tdSenderAddress = genesisAddress
          , tdParameter = untypeValue $ toVal addressB
          , tdEntrypoint = DefEpName
          , tdAmount = unsafeMkMutez 100
          }

      transfer txData addressA

      -- 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.
      validate
        $ Right
        $ expectStorageConst addressA
        $ untypeValue
        $ toVal [unsafeMkMutez 75, unsafeMkMutez 100]