-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Tests for Lorentz 'UStore'.
module Test.Lorentz.UStore.Behaviour
  ( test_Roundtrip
  , test_Conversions
  , test_Script
  ) where

import qualified Data.Map as M
import Hedgehog (Gen)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Test.HUnit (Assertion, assertFailure, (@?=))
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (testCase)

import Lorentz.Base
import Lorentz.Instr as L
import Lorentz.Pack
import Lorentz.Run
import Lorentz.UStore
import Lorentz.Value
import Lorentz.UStore.Types (genUStoreFieldExt, genUStoreSubMap)
import Michelson.Test.Dummy
import Michelson.Test.Util

data MyTemplate = MyTemplate
  { ints :: Integer |~> ()
  , bool :: UStoreField Bool
  } deriving stock (Eq, Show, Generic)

genMyTemplate :: Gen MyTemplate
genMyTemplate = MyTemplate
  <$> genUStoreSubMap (Gen.integral (Range.linearFrom 0 -1000 1000)) (pure ())
  <*> genUStoreFieldExt Gen.bool

data MyTemplateBig = MyTemplateBig
  { small :: MyTemplate
  , bytes :: ByteString |~> Natural
  , total :: UStoreField Integer
  } deriving stock (Eq, Show, Generic)

genMyTemplateBig :: Gen MyTemplateBig
genMyTemplateBig = MyTemplateBig
  <$> genMyTemplate
  <*> genUStoreSubMap (Gen.bytes (Range.linear 0 100)) (Gen.integral (Range.linear 0 1000))
  <*> genUStoreFieldExt (Gen.integral (Range.linearFrom 0 -1000 1000))

data MyMarker :: UStoreMarkerType
instance KnownUStoreMarker MyMarker where
  mkFieldMarkerUKey name = lPackValue ([mt|X|] <> name)

data MyTemplateWithMarker = MyTemplateWithMarker
  { mint :: UStoreField Integer
  , mbool :: UStoreFieldExt MyMarker Bool
  } deriving stock (Eq, Show, Generic)

genMyTemplateWithMarker :: Gen MyTemplateWithMarker
genMyTemplateWithMarker = MyTemplateWithMarker
  <$> genUStoreFieldExt (Gen.integral (Range.linearFrom 0 -1000 1000))
  <*> genUStoreFieldExt Gen.bool

test_Roundtrip :: [TestTree]
test_Roundtrip =
  [ roundtripTree genMyTemplate (mkUStore @MyTemplate) ustoreDecomposeFull
  , roundtripTree genMyTemplateBig (mkUStore @MyTemplateBig) ustoreDecomposeFull
  , roundtripTree genMyTemplateWithMarker (mkUStore @MyTemplateWithMarker) ustoreDecomposeFull
  ]

test_Conversions :: [TestTree]
test_Conversions =
  [ testGroup "Simple store template"
    [ testCase "No action" $
        ustoreChangeTest
          ( nop
          , MyTemplate (UStoreSubMap def) (UStoreField False)
          , MyTemplate (UStoreSubMap def) (UStoreField False)
          )
    , testCase "Insert into submap" $
        ustoreChangeTest
          ( unit # push 5 # ustoreInsert #ints
          , MyTemplate (UStoreSubMap def) (UStoreField False)
          , MyTemplate (UStoreSubMap $ one (5, ())) (UStoreField False)
          )
    , testCase "Delete from submap" $
        ustoreChangeTest
          ( push 3 # ustoreDelete #ints
          , MyTemplate (UStoreSubMap $ one (3, ())) (UStoreField False)
          , MyTemplate (UStoreSubMap mempty) (UStoreField False)
          )
    , testCase "Get from submap" $
        ustoreChangeTest
          ( dup # push 0 # ustoreGet #ints #
            ifNone (push 10) (L.drop # push 11) # dip unit # ustoreInsert #ints
          , MyTemplate (UStoreSubMap $ one (0, ())) (UStoreField False)
          , MyTemplate (UStoreSubMap $ M.fromList [(0, ()), (11, ())]) (UStoreField False)
          )
    , testCase "Set field" $
        ustoreChangeTest
          ( push True # ustoreSetField #bool
          , MyTemplate (UStoreSubMap mempty) (UStoreField False)
          , MyTemplate (UStoreSubMap mempty) (UStoreField True)
          )
    , testCase "Get field" $
        ustoreChangeTest
          ( ustoreGetField #bool #
            if_ (push 5) (push 0) # dip unit # ustoreInsert #ints
          , MyTemplate (UStoreSubMap mempty) (UStoreField False)
          , MyTemplate (UStoreSubMap $ one (0, ())) (UStoreField False)
          )
    , testCase "Leave some entries untouched" $
        ustoreChangeTest
          ( push 0 # ustoreDelete #ints #
            unit # push 2 # ustoreInsert #ints
          , MyTemplate (UStoreSubMap $ M.fromList [(0, ()), (1, ())]) (UStoreField False)
          , MyTemplate (UStoreSubMap $ M.fromList [(1, ()), (2, ())]) (UStoreField False)
          )
    ]

  , testGroup "Non-flat store template"
    [ testCase "Custom scenario 1" $
        ustoreChangeTest
          ( push "a" # ustoreDelete #bytes #
            push 2 # push "b" # ustoreInsert #bytes #
            ustoreGetField #total # push @Integer 1 # add # ustoreSetField #total #
            unliftUStore #small #
            unit # push 0 # ustoreInsert #ints #
            push True # ustoreSetField #bool #
            liftUStore #small
          , MyTemplateBig
            { small = MyTemplate (UStoreSubMap def) (UStoreField False)
            , bytes = UStoreSubMap $ one ("a", 1)
            , total = UStoreField 10
            }
          , MyTemplateBig
            { small = MyTemplate (UStoreSubMap $ one (0, ())) (UStoreField True)
            , bytes = UStoreSubMap $ one ("b", 2)
            , total = UStoreField 11
            }
          )
    ]
  ]
  where
    -- We accept a tuple as argument to avoid many parentheses
    ustoreChangeTest
      :: ( Each [Eq, Show, Generic] '[template]
         , UStoreTraversable MkUStoreTW template
         , UStoreTraversable DecomposeUStoreTW template
         , HasCallStack
         )
      => ( '[UStore template] :-> '[UStore template]
         , template
         , template
         )
      -> Assertion
    ustoreChangeTest (instr, initStoreHs, expectedNewStore) =
      let
        initStore = mkUStore initStoreHs
        ustore =
          leftToPrettyPanic $
          interpretLorentzLambda dummyContractEnv instr initStore
      in case ustoreDecomposeFull ustore of
          Left err -> assertFailure (toString err)
          Right ustoreHs -> ustoreHs @?= expectedNewStore

test_Script :: [TestTree]
test_Script =
  [ testCase "Only fields" $
      ustoreScriptTest MyTemplate
        { ints = UStoreSubMap mempty
        , bool = UStoreField True
        }

  , testCase "Fields and submaps" $
      ustoreScriptTest MyTemplate
        { ints = UStoreSubMap $ one (5, ())
        , bool = UStoreField True
        }

  , testCase "Complex" $
      ustoreScriptTest MyTemplateBig
        { small = MyTemplate (UStoreSubMap $ one (0, ())) (UStoreField True)
        , bytes = UStoreSubMap $ one ("b", 2)
        , total = UStoreField 11
        }
  ]
  where
    ustoreScriptTest
      :: ( Each [Eq, Show, Generic] '[template]
         , UStoreTraversable FillUStoreTW template
         , UStoreTraversable DecomposeUStoreTW template
         , HasCallStack
         )
      => template
      -> Assertion
    ustoreScriptTest store =
      let
        filling = migrationToLambda (fillUStore store)
        ustoreFilled =
          leftToPrettyPanic $
          interpretLorentzLambda dummyContractEnv filling (mkUStore ())
      in case ustoreDecomposeFull ustoreFilled of
          Left err -> assertFailure (toString err)
          Right ustoreHs -> ustoreHs @?= store