{-# LANGUAGE DeriveAnyClass #-}

-- | Tests for Lorentz compilation which uses 'LorentzCompilationWay'.
module Test.Lorentz.EntryPoints
  ( test_FieldAnnotations
  , test_TypeAnnotations
  ) where

import Test.HUnit ((@?=))
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (testCase)

import Lorentz ((:!))
import qualified Lorentz as L
import Lorentz.EntryPoints
import Lorentz.Run
import Michelson.Typed
import Michelson.Untyped (FieldAnn, TypeAnn, ann, noAnn)

data MyEntryPoints1
  = Do1 Integer
  | Do2 (Integer, Integer)
  | Do3 MyEntryPoints2
  | Do4 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 MyParams = MyParams
  { param1 :: ()
  , param2 :: ByteString
  }
  deriving stock Generic
  deriving anyclass IsoValue

instance ParameterEntryPoints MyEntryPoints1 where
  parameterEntryPoints = pepRecursive

newtype MyEntryPoints1' = MyEntryPoints1' MyEntryPoints1
  deriving newtype IsoValue

instance ParameterEntryPoints MyEntryPoints1' where
  parameterEntryPoints = mapParameterEntryPoints MyEntryPoints1' pepPlain

instance ParameterEntryPoints MyEntryPoints2 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints3 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints4 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints5 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints6 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints7 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints8 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints9 where
  parameterEntryPoints = pepPlain

instance ParameterEntryPoints MyEntryPoints10 where
  parameterEntryPoints = pepPlain

contract1 :: L.Contract MyEntryPoints1 ()
contract1 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract1' :: L.Contract MyEntryPoints1' ()
contract1' = L.drop L.# L.unit L.# L.nil L.# L.pair

contract2 :: L.Contract MyEntryPoints2 ()
contract2 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract3 :: L.Contract MyEntryPoints3 ()
contract3 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract4 :: L.Contract MyEntryPoints4 ()
contract4 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract5 :: L.Contract MyEntryPoints5 ()
contract5 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract6 :: L.Contract MyEntryPoints6 ()
contract6 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract7 :: L.Contract MyEntryPoints7 ()
contract7 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract8 :: L.Contract MyEntryPoints8 ()
contract8 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract9 :: L.Contract MyEntryPoints9 ()
contract9 = L.drop L.# L.unit L.# L.nil L.# L.pair

contract10 :: L.Contract MyEntryPoints10 ()
contract10 = 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 instance Eq (FieldAnnTree t)
deriving 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 instance Eq (TypeAnnTree t)
deriving 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 contract2)
      @?=
      Node
        (ann "do10") Leaf
        (ann "do11") Leaf

  , testGroup "Complex parameter"
    [ testCase "Interpreting as direct list of entrypoints" $
        (paramAnnTree $ compileLorentzContract contract1')
        @?=
        Node
          noAnn (Node (ann "do1") Leaf (ann "do2") Leaf)
          noAnn (Node (ann "do3") (Node noAnn Leaf noAnn Leaf) (ann "do4") Leaf)

    , testCase "Recursive entrypoints traversal" $
        (paramAnnTree $ compileLorentzContract contract1)
        @?=
        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 contract3)
      @?=
      (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 contract4)
      @?=
      (TANodeOr noAnn
        (TANodePair (ann "viewarg1") (TALeaf (ann "owner")) (TALeaf noAnn)) (TALeaf noAnn))

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

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

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

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

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

  , testCase "Big map field parameter" $
      (paramAnnTree $ compileLorentzContract contract10)
      @?=
      (TANodeOr noAnn
        (TANodeLambda (ann "bigmaparg") (TANodeBigMap noAnn (TANodePair noAnn (TALeaf (ann "balance")) (TALeaf (ann "address")))) (TALeaf noAnn)) (TALeaf noAnn))
  ]
  where
    paramAnnTree :: FullContract cp st -> TypeAnnTree cp
    paramAnnTree = extractTypeAnnTres . fcParamNotes