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

-- | Tests on Lorentz storage annotation parsing.

module Test.Lorentz.Storage
  ( test_FieldAnnotations
  , test_TypeAnnotations
  ) where

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

import Lorentz ((:!))
import Lorentz qualified as L
import Lorentz.Annotation (HasAnnotation)
import Lorentz.Run hiding (Contract(..))
import Lorentz.Value
import Morley.Michelson.Typed (Contract, Contract'(..))
import Morley.Michelson.Untyped.Annotation
import Test.Cleveland.Instances ()

import Test.Util.Annotation

----------------------------------------------------------------------------
-- Storage declarations
----------------------------------------------------------------------------

data MyParams = MyParams
  { p1 :: Natural
  , p2 :: Address
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data UnitStorage = UnitStorage
  { unit :: ()
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage1 = MyStorage1
  { st1 :: ()
  , st2 :: ()
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage2 = MyStorage2
  { st3 :: Maybe Address
  , st4 :: (Natural, Natural)
  , st5 :: MyParams
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage2r = MyStorage2r
  { st3r :: Maybe Address
  , st4r :: (Natural, Natural)
  , st5r :: MyParams
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage3 = MyStorage3
  { st6 :: ("tuplearg" :! ("TL" :! Address, "TR" :! Integer), "boolarg" :! Bool)
  , st7 :: ("integerarg" :! Natural, "boolarg" :! Bool)
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage4 = MyStorage4
  { st8 :: ("bigmaparg" :! L.Lambda (BigMap Natural ("balance" :! Natural , "address" :! L.Address)) ())
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage5 = MyStorage5
  { st9 :: ("maybearg" :! Maybe ("maybeinner" :! Natural))
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage6 = MyStorage6
  { st10 :: ("lambdaarg" :! L.Lambda Natural Natural)
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

data MyStorage7 = MyStorage7
  { st11 :: ("listarg" :! [("balance" :! Natural , "address" :! L.Address)])
  } deriving stock Generic
    deriving anyclass (IsoValue, HasAnnotation)

dummyContract :: forall storage. L.NiceStorageFull storage => L.Contract () storage ()
dummyContract = defaultContract L.fakeCoerce

----------------------------------------------------------------------------
-- Field annotations
----------------------------------------------------------------------------

test_FieldAnnotations :: [TestTree]
test_FieldAnnotations =
  [
    testCase "Simple storage" $ do
      extractAnnotation @MyStorage1
      @?=
      FANodePair
        "st1" FALeaf
        "st2" FALeaf

 , testCase "Complex storage" $ do
      extractAnnotation @MyStorage2
      @?=
        FANodePair
          "st3"
            FALeaf
          noAnn
            (FANodePair
              "st4" FALeaf
              "st5"
                (FANodePair
                  "p1" FALeaf
                  "p2" FALeaf))

  , testCase "Complex parameter recursive" $ do
      extractAnnotation @MyStorage2r
      @?=
        FANodePair
          "st3r"
            FALeaf
          noAnn
            (FANodePair
              "st4r"
                FALeaf
              "st5r"
                (FANodePair
                  "p1" FALeaf
                  "p2" FALeaf))
  ]
  where
    storageAnnTree :: Contract cp st -> FieldAnnTree st
    storageAnnTree = extractFieldAnnTree . cStoreNotes

    extractAnnotation :: forall st. L.NiceStorageFull st => FieldAnnTree (ToT st)
    extractAnnotation = storageAnnTree $ L.toMichelsonContract (dummyContract @st)

test_TypeAnnotations :: [TestTree]
test_TypeAnnotations =
  [ testCase "Address primitive storage with no annotations" $
      extractAnnotation @Address
      @?=
      TALeaf noAnn

  , testCase "Named type annotation" $
      extractAnnotation @MyStorage3
      @?=
      TANodePair noAnn
        (TANodePair noAnn
          (TANodePair "tuplearg"
            (TALeaf "TL")
            (TALeaf "TR"))
          (TALeaf "boolarg"))
        (TANodePair noAnn
          (TALeaf "integerarg")
          (TALeaf "boolarg"))

  , testCase "BigMap type annotation" $
      extractAnnotation @MyStorage4
      @?=
      TANodeLambda "bigmaparg"
        (TANodeBigMap noAnn
          (TALeaf noAnn)
          (TANodePair noAnn
            (TALeaf "balance")
            (TALeaf "address")))
        (TALeaf noAnn)

  , testCase "Maybe type annotation" $
      extractAnnotation @MyStorage5
      @?=
      TANodeOption "maybearg" (TALeaf "maybeinner")

  , testCase "Lambda type annotation" $
      extractAnnotation @MyStorage6
      @?=
      TANodeLambda "lambdaarg" (TALeaf noAnn) (TALeaf noAnn)

  , testCase "List type annotation" $
      extractAnnotation @MyStorage7
      @?=
      TANodeList "listarg"
        (TANodePair noAnn
          (TALeaf "balance")
          (TALeaf "address"))

  ]
  where
    storageAnnTree :: Contract cp st -> TypeAnnTree st
    storageAnnTree = extractTypeAnnTree . cStoreNotes

    extractAnnotation :: forall st. L.NiceStorageFull st => TypeAnnTree (ToT st)
    extractAnnotation = storageAnnTree $ L.toMichelsonContract (dummyContract @st)