{-# LANGUAGE DeriveAnyClass #-}

-- | Tests for Lorentz compilation which uses 'LorentzCompilationWay'.
module Test.Lorentz.EntryPoints
  ( test_FieldAnnotations
  , test_TypeAnnotations
  , 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 Test.Util.TypeSpec (ExactlyIs)

import Lorentz ((:!), ( # ), (/->))
import qualified Lorentz as L
import Lorentz.Constraints
import Lorentz.EntryPoints
import Lorentz.Run
import Lorentz.Test
import Lorentz.Value
import Michelson.Typed
import Michelson.Untyped (FieldAnn, TypeAnn, ann, noAnn)

----------------------------------------------------------------------------
-- 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

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" L.:? ("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 MyEntryPointsWithDef
  = Default Integer
  | NonDefault Natural
  deriving stock Generic
  deriving anyclass IsoValue

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

newtype MyNewtypeParam a = MyNewtypeParam a
  deriving stock Generic
  deriving anyclass IsoValue

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 MyEntryPointsWithDef where
  type ParameterEntryPointsDerivation MyEntryPointsWithDef = EpdPlain

instance ParameterHasEntryPoints (MyNewtypeParam Integer) where
  type ParameterEntryPointsDerivation (MyNewtypeParam Integer) = EpdPlain

dummyContract :: L.Contract param ()
dummyContract = L.drop L.# L.unit L.# L.nil L.# L.pair

-- | Helper datatype which contains field annotations from 'NTOr'.
data FieldAnnTree t where
  Leaf :: FieldAnnTree t
  Node :: FieldAnn -> FieldAnnTree a
       -> FieldAnn -> FieldAnnTree b
       -> FieldAnnTree ('TOr a b)

deriving stock instance Eq (FieldAnnTree t)
deriving stock instance Show (FieldAnnTree t)

data TypeAnnTree t where
  TALeaf :: TypeAnn -> TypeAnnTree t
  TANodeOption :: TypeAnn -> TypeAnnTree a -> TypeAnnTree ('TOption a)
  TANodePair :: TypeAnn -> TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TPair a b)
  TANodeOr :: TypeAnn -> TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TOr a b)
  TANodeLambda :: TypeAnn -> TypeAnnTree a -> TypeAnnTree b -> TypeAnnTree ('TLambda a b)
  TANodeList :: TypeAnn -> TypeAnnTree a -> TypeAnnTree ('TList a)
  TANodeMap :: TypeAnn -> TypeAnnTree b -> TypeAnnTree ('TMap a b)
  TANodeBigMap :: TypeAnn -> TypeAnnTree b -> TypeAnnTree ('TBigMap a b)

deriving stock instance Eq (TypeAnnTree t)
deriving stock instance Show (TypeAnnTree t)

extractTypeAnnTres :: Notes t -> TypeAnnTree t
extractTypeAnnTres = \case
  NTc ta -> TALeaf ta
  NTKey ta -> TALeaf ta
  NTUnit ta -> TALeaf ta
  NTSignature ta -> TALeaf ta
  NTOption ta n1 -> TANodeOption ta (extractTypeAnnTres n1)
  NTList ta n1 -> TANodeList ta (extractTypeAnnTres n1)
  NTSet ta _ -> TALeaf ta
  NTOperation ta -> TALeaf ta
  NTContract ta _ -> TALeaf ta
  NTPair ta _ _ n1 n2 -> TANodePair ta (extractTypeAnnTres n1) (extractTypeAnnTres n2)
  NTOr ta _ _ n1 n2 -> TANodeOr ta (extractTypeAnnTres n1) (extractTypeAnnTres n2)
  NTLambda ta n1 n2 -> TANodeLambda ta (extractTypeAnnTres n1) (extractTypeAnnTres n2)
  NTMap ta _ n1 -> TANodeMap ta (extractTypeAnnTres n1)
  NTBigMap ta _ n1 -> TANodeBigMap ta (extractTypeAnnTres n1)
  NTChainId ta -> TALeaf ta

extractFieldAnnTree :: Notes t -> FieldAnnTree t
extractFieldAnnTree = \case
  NTOr _ lann rann lnotes rnotes ->
    Node lann (extractFieldAnnTree lnotes) rann (extractFieldAnnTree rnotes)
  _ ->
    Leaf

test_FieldAnnotations :: [TestTree]
test_FieldAnnotations =
  [ testCase "Simple parameter" $
      (paramAnnTree $ compileLorentzContract (dummyContract @MyEntryPoints2))
      @?=
      Node
        (ann "do10") Leaf
        (ann "do11") Leaf

  , testGroup "Complex parameter"
    [ testCase "Interpreting as direct list of entrypoints" $
        (paramAnnTree $ compileLorentzContract (dummyContract @MyEntryPoints1a))
        @?=
        Node
          noAnn (Node (ann "do1a") Leaf (ann "do2a") Leaf)
          noAnn (Node (ann "do3a") (Node noAnn Leaf noAnn Leaf) (ann "do4a") Leaf)

    , testCase "Recursive entrypoints traversal" $
        (paramAnnTree $ compileLorentzContract (dummyContract @MyEntryPoints1))
        @?=
        Node
          noAnn (Node (ann "do1") Leaf (ann "do2") Leaf)
          noAnn (Node
                  noAnn (Node (ann "do10") Leaf (ann "do11") Leaf)
                  (ann "do4") Leaf
                )
    ]
  ]
  where
    paramAnnTree :: FullContract cp st -> FieldAnnTree cp
    paramAnnTree = extractFieldAnnTree . fcParamNotes

test_TypeAnnotations :: [TestTree]
test_TypeAnnotations =
  [ testCase "Named field parameter" $
      (paramAnnTree $ compileLorentzContract (dummyContract @MyEntryPoints3))
      @?=
      (TANodeOr noAnn
       (TANodePair noAnn (TANodePair (ann "tuplearg") (TALeaf (ann "TL")) (TALeaf (ann "TR"))) (TALeaf (ann "boolarg")))
       (TANodePair noAnn (TALeaf (ann "integerarg")) (TALeaf (ann "boolarg")))
      )
  , testCase "Named field parameter for views" $
      (paramAnnTree $ compileLorentzContract (dummyContract @MyEntryPoints4))
      @?=
      (TANodeOr noAnn
        (TANodePair (ann "viewarg1") (TALeaf (ann "owner")) (TALeaf noAnn)) (TALeaf noAnn))

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

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

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

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

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

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

  , testCase "Newtype" $
      (paramAnnTree $ compileLorentzContract (dummyContract @(MyNewtypeParam Integer)))
      @?=
      TALeaf noAnn

  , testGroup "Primitive type parameter"
      -- Parameters used in these test cases should not require any instances
      [ testCase "Address" $
          (paramAnnTree $ compileLorentzContract (dummyContract @Address))
          @?=
          TALeaf noAnn
      , testCase "Void" $
          (paramAnnTree $ compileLorentzContract (dummyContract @(L.Void_ Integer Natural)))
          @?=
          TANodePair noAnn (TALeaf noAnn) (TANodeLambda noAnn (TALeaf noAnn) (TALeaf noAnn))
      ]
  ]
  where
    paramAnnTree :: FullContract cp st -> TypeAnnTree cp
    paramAnnTree = extractTypeAnnTres . fcParamNotes

----------------------------------------------------------------------------
-- @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)
    ]
  ]

-- | 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
  -> L.Contract Address ()
callerContract epRef arg =
  L.car # L.contractCalling @cp epRef #
  L.assertSome [mt|Contract lookup failed|] #
  L.push (toMutez 1) # L.push arg # L.transferTokens #
  L.dip (L.unit # L.nil) # L.cons # L.pair

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

        caller <- lOriginate myCallerContract "Caller" () (toMutez 10)
        target <- lOriginateEmpty myTargetContract "Call target"
        lCallDef caller (toAddress target)
        validate . Right $
          lExpectStorageConst target (5 :: Natural)

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

        caller <- lOriginate myCallerContract "Caller" () (toMutez 10)
        target <- lOriginateEmpty myTargetContract "Call target"
        lCallDef caller (toAddress target)
        validate . Right $
          lExpectStorageConst target (3 :: Natural)
  ]

test_Self_call :: [TestTree]
test_Self_call =
  [ testCase "Calling entrypoint" $
      integrationalTestExpectation $ do
        let myContract = L.car # L.caseT @MyEntryPoints2
              ( #cDo10 /->
                  L.selfCalling @MyEntryPoints2 (Call @"Do11") #
                  L.push (toMutez 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 <- lOriginate myContract "Contract" 0 (toMutez 10)
        lCallDef contractRef Do10
        validate . Right $
          lExpectStorageConst contractRef (15 :: Natural)
  ]