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

-- | Tests for Lorentz compilation which uses 'LorentzCompilationWay'.
module Test.Lorentz.Entrypoints
  ( test_FieldEntrypoints
  , test_TypeEntrypoints
  , test_RootEntrypoints
  , test_Entrypoints_lookup
  , test_Contract_call
  , test_Self_call
  ) where


import Fcf (Eval)
import Test.HUnit ((@?=))
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (testCase)
import Test.Tasty.TypeSpec (typeTest)
import Test.TypeSpec (Is, TypeSpec(..))

import Lorentz ((#), (/->))
import Lorentz qualified as L
import Lorentz.Annotation
import Lorentz.Constraints
import Lorentz.Entrypoints
import Lorentz.Run
import Lorentz.Value
import Morley.Michelson.Typed hiding (Contract)
import Morley.Michelson.Untyped (RootAnn, noAnn)
import Morley.Util.Named
import Test.Cleveland

import Test.Util.Annotation
import Test.Util.TypeSpec (ExactlyIs)

----------------------------------------------------------------------------
-- Entrypoints declarations
----------------------------------------------------------------------------

data MyEntrypoints1
  = Do1 Integer
  | Do2 (Integer, Integer)
  | Do3 MyEntrypoints2
  | Do4 MyParams
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints1a
  = Do1a Integer
  | Do2a (Integer, Integer)
  | Do3a MyEntrypoints2
  | Do4a MyParams
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints2
  = Do10
  | Do11 Natural
  deriving stock Generic
  deriving anyclass (IsoValue, HasAnnotation)

data MyEntrypoints3
  = Do12 ("tuplearg" :! ("TL" :! Integer, "TR" :!  Natural), "boolarg" :! Bool)
  | Do13 ("integerarg" :! Integer, "boolarg" :! Bool)
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints4
  = Do14 ("viewarg1" :! L.View_ ("owner" :! L.Address) Natural)
  | Do15 ()
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints5
  = Do16 ("maybearg" :! Maybe ("maybeinner" :! Natural))
  | Do17 ()
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints6
  = Do18 ("lambdaarg" :! L.Lambda Natural Natural)
  | Do19 ()
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints7
  = Do20 ("listarg" :! [("balance" :! Natural , "address" :! L.Address)])
  | Do21 ()
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints8
  = Do22 ("maparg" :! (Map Natural ("balance" :! Natural , "address" :! L.Address)))
  | Do23 ()
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints9
  = Do24 ("maybearg" :? ("maybeinner" :! Natural))
  | Do25 ()
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints10
  = Do26 ("bigmaparg" :! L.Lambda (BigMap Natural ("balance" :! Natural , "address" :! L.Address)) ())
  | Do27 ()
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypoints11
  = Do28 ("kek" :! Natural, "pek" :! Integer)
  | Do29
  deriving stock Generic
  deriving anyclass (IsoValue, HasAnnotation)

data MyEntrypointsWithDef
  = Default Integer
  | NonDefault Natural
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypointsWithRoot
  = Dor1
  | Dor2 Natural
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypointsWithCustomAnn
  = Doc1 Natural
  | Doc2 CustomProd
  deriving stock Generic
  deriving anyclass (IsoValue)

data CustomProd = CustomProd
  { cp1 :: Natural
  , cp2 :: CustomSum
  }
  deriving stock Generic
  deriving anyclass (IsoValue)

instance HasAnnotation CustomProd where
  annOptions = defaultAnnOptions {fieldAnnModifier = \a -> "prod_" <> a}

data CustomSum = Cs1 | Cs2
  deriving stock Generic
  deriving anyclass (IsoValue)

instance HasAnnotation CustomSum where
  annOptions = defaultAnnOptions {fieldAnnModifier = \a -> "sum_" <> a}

data MyParams = MyParams
  { param1 :: ()
  , param2 :: ByteString
  }
  deriving stock Generic
  deriving anyclass (IsoValue, HasAnnotation)

-- Normally this cannot declare entrypoints because this is not a sum type.
-- But we will declare them forcibly
data MySingleEntrypoint = Dos1 Integer
  deriving stock Generic
  deriving anyclass (IsoValue, HasAnnotation)

data MyEntrypointsDelegated
  = Dod1
  | Dod2 MyEntrypointsSubDelegated
  deriving stock Generic
  deriving anyclass (IsoValue)

data MyEntrypointsSubDelegated
  = Dosd1
  | Dosd2
  deriving stock Generic
  deriving anyclass (IsoValue, HasAnnotation)

instance ParameterHasEntrypoints MyEntrypoints1 where
  type ParameterEntrypointsDerivation MyEntrypoints1 = EpdRecursive

instance ParameterHasEntrypoints MyEntrypoints1a where
  type ParameterEntrypointsDerivation MyEntrypoints1a = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints2 where
  type ParameterEntrypointsDerivation MyEntrypoints2 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints3 where
  type ParameterEntrypointsDerivation MyEntrypoints3 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints4 where
  type ParameterEntrypointsDerivation MyEntrypoints4 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints5 where
  type ParameterEntrypointsDerivation MyEntrypoints5 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints6 where
  type ParameterEntrypointsDerivation MyEntrypoints6 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints7 where
  type ParameterEntrypointsDerivation MyEntrypoints7 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints8 where
  type ParameterEntrypointsDerivation MyEntrypoints8 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints9 where
  type ParameterEntrypointsDerivation MyEntrypoints9 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints10 where
  type ParameterEntrypointsDerivation MyEntrypoints10 = EpdPlain

instance ParameterHasEntrypoints MyEntrypoints11 where
  type ParameterEntrypointsDerivation MyEntrypoints11 = EpdNone

instance ParameterHasEntrypoints MyEntrypointsWithDef where
  type ParameterEntrypointsDerivation MyEntrypointsWithDef = EpdPlain

instance ParameterHasEntrypoints (ShouldHaveEntrypoints MySingleEntrypoint) where
  type ParameterEntrypointsDerivation (ShouldHaveEntrypoints MySingleEntrypoint) = EpdPlain

instance ParameterHasEntrypoints MyEntrypointsDelegated where
  type ParameterEntrypointsDerivation MyEntrypointsDelegated = EpdDelegate

instance ParameterHasEntrypoints MyEntrypointsSubDelegated where
  type ParameterEntrypointsDerivation MyEntrypointsSubDelegated = EpdNone

instance ParameterHasEntrypoints MyEntrypointsWithRoot where
  type ParameterEntrypointsDerivation MyEntrypointsWithRoot = EpdWithRoot "root" EpdPlain

instance ParameterHasEntrypoints MyEntrypointsWithCustomAnn where
  type ParameterEntrypointsDerivation MyEntrypointsWithCustomAnn = EpdPlain

dummyContract :: NiceParameterFull param => Contract param () ()
dummyContract = L.defaultContract $
  L.drop L.# L.unit L.# L.nil L.# L.pair

test_FieldEntrypoints :: [TestTree]
test_FieldEntrypoints =
  [ testCase "Simple parameter" $
      (paramAnnTree $ dummyContract @MyEntrypoints2)
      @?=
      FANodeOr
        "do10" FALeaf
        "do11" FALeaf

  , testGroup "Complex parameter"
    [ testCase "Interpreting as direct list of entrypoints" $
        (paramAnnTree $ dummyContract @MyEntrypoints1a)
        @?=
        FANodeOr
          noAnn (FANodeOr "do1a" FALeaf "do2a" FALeaf)
          noAnn
            (FANodeOr
              "do3a"
                (FANodeOr noAnn FALeaf noAnn FALeaf)
              "do4a"
                (FANodePair "param1" FALeaf "param2" FALeaf))

    , testCase "Recursive entrypoints traversal" $
        (paramAnnTree $ dummyContract @MyEntrypoints1)
        @?=
        FANodeOr
          noAnn (FANodeOr "do1" FALeaf "do2" FALeaf)
          noAnn (FANodeOr
                  noAnn
                    (FANodeOr "do10" FALeaf "do11" FALeaf)
                  "do4"
                    (FANodePair
                      "param1" FALeaf
                      "param2" FALeaf)
                )

    , testCase "Delegating entrypoints traversal" $
        (paramAnnTree $ dummyContract @MyEntrypointsDelegated)
        @?=
        FANodeOr "dod1" FALeaf
             "dod2" (FANodeOr noAnn FALeaf noAnn FALeaf)
    , testCase "Parameter with custom fieldAnnotation" $
        (paramAnnTree $ dummyContract @MyEntrypointsWithCustomAnn)
        @?=
        FANodeOr
          "doc1" FALeaf
          "doc2"
          (FANodePair
            "prod_cp1" FALeaf
            "prod_cp2"
            (FANodeOr
              "sum_cs1" FALeaf
              "sum_cs2" FALeaf
            )
          )
    ]
  ]
  where
    paramAnnTree :: Contract cp st () -> FieldAnnTree (ToT cp)
    paramAnnTree = extractFieldAnnTree . pnNotes . cParamNotes . L.toMichelsonContract

test_TypeEntrypoints :: [TestTree]
test_TypeEntrypoints =
  [ testCase "Named field parameter" $
      (paramAnnTree $ dummyContract @MyEntrypoints3)
      @?=
      (TANodeOr noAnn
       (TANodePair noAnn (TANodePair "tuplearg" (TALeaf "TL") (TALeaf "TR")) (TALeaf "boolarg"))
       (TANodePair noAnn (TALeaf "integerarg") (TALeaf "boolarg"))
      )
  , testCase "Named field parameter for views" $
      (paramAnnTree $ dummyContract @MyEntrypoints4)
      @?=
      (TANodeOr noAnn
        (TANodePair "viewarg1" (TALeaf "owner") (TALeaf noAnn)) (TALeaf noAnn))

  , testCase "Maybe field parameter" $
      (paramAnnTree $ dummyContract @MyEntrypoints5)
      @?=
      (TANodeOr noAnn
        (TANodeOption "maybearg" (TALeaf "maybeinner")) (TALeaf noAnn))

  , testCase "Lambda field parameter" $
      (paramAnnTree $ dummyContract @MyEntrypoints6)
      @?=
      (TANodeOr noAnn
        (TANodeLambda "lambdaarg" (TALeaf noAnn) (TALeaf noAnn)) (TALeaf noAnn))

  , testCase "List field parameter" $
      (paramAnnTree $ dummyContract @MyEntrypoints7)
      @?=
      (TANodeOr noAnn
       (TANodeList "listarg" (TANodePair noAnn (TALeaf "balance") (TALeaf "address"))) (TALeaf noAnn))

  , testCase "Map field parameter" $
      (paramAnnTree $ dummyContract @MyEntrypoints8)
      @?=
      (TANodeOr noAnn
        (TANodeMap
          "maparg"
          (TALeaf noAnn)
          (TANodePair noAnn (TALeaf "balance") (TALeaf "address"))
        )
        (TALeaf noAnn)
      )

  , testCase "Maybe field parameter 2" $
      (paramAnnTree $ dummyContract @MyEntrypoints9)
      @?=
      (TANodeOr noAnn
        (TANodeOption "maybearg" (TALeaf "maybeinner")) (TALeaf noAnn))

  , testCase "Big map field parameter" $
      (paramAnnTree $ dummyContract @MyEntrypoints10)
      @?=
      (TANodeOr noAnn
        (TANodeLambda "bigmaparg"
          (TANodeBigMap
            noAnn
            (TALeaf noAnn)
            (TANodePair noAnn (TALeaf "balance") (TALeaf "address"))
          )
          (TALeaf noAnn)
        )
        (TALeaf noAnn)
      )

  , testCase "Newtype" $
      (paramAnnTree $ dummyContract @(ShouldHaveEntrypoints MySingleEntrypoint))
      @?=
      TALeaf noAnn

  , testGroup "Primitive type parameter"
      -- Parameters used in these test cases should not require any instances
      [ testCase "Address" $
          (paramAnnTree $ dummyContract @Address)
          @?=
          TALeaf noAnn
      , testCase "Void" $
          (paramAnnTree $ dummyContract @(L.Void_ Integer Natural))
          @?=
          TANodePair noAnn (TALeaf noAnn) (TANodeLambda noAnn (TALeaf noAnn) (TALeaf noAnn))
      ]
  , testCase "EpdNone case (type annotations are preserved)" $
      (paramAnnTree $ dummyContract @MyEntrypoints11)
      @?=
      (TANodeOr noAnn
        (TANodePair noAnn (TALeaf "kek") (TALeaf "pek"))
        (TALeaf noAnn)
      )
  ]
  where
    paramAnnTree :: Contract cp st () -> TypeAnnTree (ToT cp)
    paramAnnTree = extractTypeAnnTree . pnNotes . cParamNotes . L.toMichelsonContract

test_RootEntrypoints :: [TestTree]
test_RootEntrypoints =
  [ testCase "No root entrypoint" $
      (rootAnn $ dummyContract @MyEntrypoints1)
      @?=
      noAnn

  , testCase "With root entrypoint" $
      (rootAnn $ dummyContract @MyEntrypointsWithRoot)
      @?=
      "root"
  ]
  where
    rootAnn :: Contract cp st () -> RootAnn
    rootAnn = pnRootAnn . cParamNotes . L.toMichelsonContract

----------------------------------------------------------------------------
-- @contract@ instruction
----------------------------------------------------------------------------

test_Entrypoints_lookup :: [TestTree]
test_Entrypoints_lookup =
  [ testGroup "Flat parameter type"
    [ typeTest "Default entrypoint arg" $ Valid @
        (GetDefaultEntrypointArg MyEntrypoints1a `Is` MyEntrypoints1a)
    , typeTest "Can get entrypoint on surface" $ Valid @
        (GetEntrypointArg MyEntrypoints1a "Do1a" `Is` Integer)
    , typeTest "Cannot get entrypoint in deep" $ Valid @
        (Eval (LookupParameterEntrypoint MyEntrypoints1a "Do11")
           `ExactlyIs` 'Nothing
        )
    ]
  , testGroup "Nested parameter type"
    [ typeTest "Default entrypoint arg" $ Valid @
        (GetDefaultEntrypointArg MyEntrypoints1 `Is` MyEntrypoints1)
    , typeTest "Can get entrypoint on surface" $ Valid @
        (GetEntrypointArg MyEntrypoints1 "Do1" `Is` Integer)
    , typeTest "Can get entrypoint in deep" $ Valid @
        (GetEntrypointArg MyEntrypoints1 "Do11" `Is` Natural)
    , typeTest "Can get entrypoint without arg" $ Valid @
        (GetEntrypointArg MyEntrypoints1 "Do10" `Is` ())
    ]
  , testGroup "Parameter type with default entrypoint"
    [ typeTest "Default entrypoint arg" $ Valid @
        (GetDefaultEntrypointArg MyEntrypointsWithDef `Is` Integer)
    , typeTest "Can get non-default entrypoint" $ Valid @
        (GetEntrypointArg MyEntrypointsWithDef "NonDefault" `Is` Natural)
    ]
  , testGroup "Delegation"
    [ typeTest "Calling immediate entrypoint works" $ Valid @
        (GetEntrypointArg MyEntrypointsDelegated "Dod1" `Is` ())
    ]
  ]

-- | A contract which accepts 'Address' as parameter and calls specific
-- entrypoint of another contract.
callerContract
  :: forall cp mname arg.
     ( arg ~ GetEntrypointArgCustom cp mname
     , NiceConstant arg, NiceParameter arg, NiceParameterFull cp
     )
  => EntrypointRef mname
  -> arg
  -> Contract (TAddress cp ()) () ()
callerContract epRef argument = L.defaultContract $
  L.car # L.contractCalling epRef #
  L.assertSome ("Contract lookup failed" :: MText) #
  L.push 1 # L.push argument # L.transferTokens #
  L.dip (L.unit # L.nil) # L.cons # L.pair

test_Contract_call :: [TestTree]
test_Contract_call =
  [ testScenario "Calling entrypoint" $ scenario do
      let myCallerContract = callerContract (Call @"Do11") 5
      let myTargetContract = L.defaultContract $ L.car # L.caseT @MyEntrypoints2
            ( #cDo10 /-> L.push 0
            , #cDo11 /-> L.nop
            ) # L.nil # L.pair

      caller <- originate $ OriginateData "Caller" 10 () myCallerContract
      target <- originateSimple "Call target" def myTargetContract
      call caller CallDefault (toTAddress target)
      getStorage target @@== 5

  , testScenario "Calling default entrypoint" $ scenario do
      let myCallerContract = callerContract CallDefault 3
      let myTargetContract = L.defaultContract $ L.car # L.caseT @MyEntrypointsWithDef
            ( #cDefault /-> L.nop
            , #cNonDefault /-> L.neg
            ) # L.nil # L.pair

      caller <- originate$ OriginateData "Caller" 10 () myCallerContract
      target <- originateSimple "Call target" def myTargetContract
      call caller CallDefault (toTAddress target)
      getStorage target @@== 3

  , testScenario "Calling root entrypoint" $ scenario do
      let myCallerContract = callerContract (Call @"root") (Dor2 5)
      let myTargetContract = L.defaultContract $ L.car # L.caseT @MyEntrypointsWithRoot
            ( #cDor1 /-> L.push 0
            , #cDor2 /-> L.nop
            ) # L.nil # L.pair

      caller <- originate $ OriginateData "Caller" 10 () myCallerContract
      target <- originateSimple "Call target" def myTargetContract
      call caller CallDefault (toTAddress target)
      getStorage target @@== 5
  ]

test_Self_call :: [TestTree]
test_Self_call =
   [ testScenario "Calling entrypoint" $ scenario do
      let myContract = mkContract $ L.car # L.caseT @MyEntrypoints2
            ( #cDo10 /->
                L.selfCalling @MyEntrypoints2 (Call @"Do11") #
                L.push 1 # L.push 5 # L.transferTokens #
                L.dip (L.push @Integer 1 # L.nil) # L.cons # L.pair
            , #cDo11 /-> L.push @Integer 10 # L.add # L.nil # L.pair
            )

      contractRef <- originate $ OriginateData "Contract" 10 0 myContract
      call contractRef CallDefault Do10
      getStorage contractRef @@== 15

  , testScenario "Calling root entrypoint" $ scenario do
      let myContract = L.defaultContract $ L.car # L.caseT @MyEntrypointsWithRoot
            ( #cDor1 /->
                L.selfCalling @MyEntrypointsWithRoot (Call @"root") #
                L.push 1 # L.push (Dor2 5) # L.transferTokens #
                L.dip (L.push @Integer 1 # L.nil) # L.cons # L.pair
            , #cDor2 /->
                L.push @Integer 10 # L.add # L.nil # L.pair
            )
      contractRef <- originate $ OriginateData "Contract" 10 0 myContract
      call contractRef CallDefault Dor1
      getStorage contractRef @@== 15
  ]