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

{-# LANGUAGE OverloadedLists #-}

module TestSuite.Cleveland.StorageWithBigMaps
  ( test_EMPTY_BIG_MAP
  , test_UPDATE
  , test_DUP
  , test_UniqueIDs
  , test_getAllBigMapValues
  , test_getAllBigMapValues_ConsistentWith_getBigMapValue
  , test_getBigMapSize
  , test_GetBigMapValue
  , test_GetBigMapValueMaybe_InvalidBigMapId
  , test_GetBigMapValueMaybe_InvalidKey
  , test_ID_BecomesInvalid
  ) where

import Lorentz as L hiding (assert)

import Control.Lens (to)
import Data.Data (Data)
import Data.Data.Lens (biplate)
import Fmt (Doc, blockListF, unlinesF)
import Test.Tasty (TestTree, testGroup)

import Test.Cleveland
import TestSuite.Util (idContract, saveInStorageContract)

data BigMaps = BigMaps
  { _bmBigMap :: BigMap Natural Natural
  , _bmOpt :: Maybe (BigMap Natural Natural)
  , _bmList :: [BigMap Natural Natural]
  , _bmPair :: (BigMap Natural Natural, BigMap Natural Natural)
  , _bmLeft :: Either (BigMap Natural Natural) (BigMap Natural Natural)
  , _bmRight :: Either (BigMap Natural Natural) (BigMap Natural Natural)
  , _bmMap :: Map Natural (BigMap Natural Natural)
  }
  deriving stock (Generic)
  deriving anyclass (IsoValue, HasAnnotation)

deriveRPC "BigMaps"
deriving stock instance Data BigMapsRPC

-- | There are 5 ways of creating big_maps on the chain:
--
-- * Originate a contract with big_maps in its initial storage.
-- * Call a contract with a parameter with big_maps in it, and then save those big_maps in the contract's storage.
-- * While executing a contract's code:
--     * execute the `EMPTY_BIG_MAP` instruction.
--     * execute the `DUP` instruction to duplicate a big_map or a value containing big_map(s).
--     * execute the `DUP n` instruction to duplicate a big_map or a value containing big_map(s).
--
-- This contract does all of those things.
createBigMaps :: Contract BigMaps BigMaps ()
createBigMaps = defaultContract $
  unpair
  # L.dip (getField #_bmList)

  -- Extract all the big_maps that were passed in through the transfer's parameter
  -- and save them in the contract's storage.
  # extractBigMaps

  -- Create new big_maps via: EMPTY_BIG_MAP, DUP, DUP n.
  -- Note: instead of duplicating a `big_map k v`, we duplicate a `list (big_map k v)` to ensure
  -- the algorithm works even on deeply nested big_maps.
  # nil # emptyBigMap # cons
  # dup
  # dupN @2

  -- Merge all the big_maps into the storage's `_bmList` field
  # joinList # joinList # joinList # joinList
  # setField #_bmList

  # nil @Operation
  # pair

  where
    joinList :: List a : List a : s :-> List a : s
    joinList = L.iter cons

    extractBigMaps :: BigMaps : s :-> List (BigMap Natural Natural) : s
    extractBigMaps =
      L.listE
        [ L.getField #_bmBigMap
        , L.getField #_bmOpt # L.assertSome @MText "Expected Some, got None"
        , getField #_bmLeft # assertLeft @MText "Expected Left, got Right"
        , getField #_bmRight # assertRight @MText "Expected Right, got Left"
        ]
      # L.swap

      # getField #_bmList
      # L.swap # dip joinList

      # getField #_bmPair
      # L.swap # dip (unpair # dip cons # cons)

      # getField #_bmMap
      # L.swap # dip (iter (cdr # cons))

      # L.drop

test_EMPTY_BIG_MAP :: TestTree
test_EMPTY_BIG_MAP =
  testScenario "EMPTY_BIG_MAP always creates a new ID" $ scenario do
    forM_ contracts \c ->  do
      addr <- originate "contract" [("a", 1)] c
      id0 <- getStorage addr
      transfer addr
      unBigMapId <$> getStorage addr @@/= unBigMapId id0
  where
    contracts :: [Contract () (BigMap MText Integer) ()]
    contracts = [createEmptyBigMap, copyToEmptyBigMap]

    createEmptyBigMap :: Contract () (BigMap MText Integer) ()
    createEmptyBigMap = L.mkContractWith L.intactCompilationOptions $ L.mkContractCode $
      L.drop
      # L.emptyBigMap
      # L.nil # L.pair

    copyToEmptyBigMap :: Contract () (BigMap MText Integer) ()
    copyToEmptyBigMap = L.mkContractWith L.intactCompilationOptions $ L.mkContractCode $
      L.drop
      # L.emptyBigMap @MText @Integer
      # L.push @Integer 1
      # L.some
      # L.push @MText "a"
      # L.update
      # L.nil # L.pair

test_UPDATE :: TestTree
test_UPDATE =
  testScenario "UPDATE never creates a new ID" $ scenario do
    addr <- originate "contract" [("12", 12)] updateBigMap
    id0 <- getStorage addr
    transfer addr
    unBigMapId <$> getStorage addr @@== unBigMapId id0
  where
    updateBigMap :: Contract () (BigMap MText Integer) ()
    updateBigMap = L.mkContractWith L.intactCompilationOptions $ L.mkContractCode $
      L.cdr
      # L.push 12 # L.some
      # L.push "12"
      # L.update
      # L.nil # L.pair

test_DUP :: TestTree
test_DUP =
  -- This is a case where the emulator and the network's implementation diverge *slightly*.
  --
  --   * In the emulator, the @DUP@ instruction _always_ creates a new bigmap ID.
  --   * In a network, the @DUP@ instruction _usually_ creates a new bigmap ID,
  --     except in a very specific situation: when you duplicate a bigmap, save
  --     the duplicate in the storage and discard the old bigmap.
  --     When this happens, the old bigmap's ID is reused and assigned to the new bigmap.
  --
  -- I don't think the fact that the two implementations diverge is a problem per se,
  -- because, IMO, users shouldn't rely on bigmap's IDs remaining the same
  -- across contract calls anyways.
  --
  -- This test is here as a means of documentation.
  testGroup "DUP"
    [ testScenario "BigMap ID may be reused when old BigMap is discarded" $ scenario do
        addr <- originate "contract" [("a", 1)] dupAndDiscardOld

        originalId <- getStorage addr
        transfer addr
        dupId <- getStorage addr

        ifEmulation @()
          -- In the emulator, the dupped bigmap is always assigned a brand new ID.
          (do unBigMapId dupId @/= unBigMapId originalId)
          -- In a network, the old ID might be reused.
          (do unBigMapId dupId @== unBigMapId originalId)

    , testScenario "A new BigMap ID is generated when the old BigMap is not discarded" $ scenario do
        addr <- originate "contract" ([("a", 1)], [("a", 1)]) dupAndKeepBoth

        (id0, _) <- getStorage addr
        transfer addr
        (id1, dupId) <- getStorage addr

        -- The original bigmap's ID did not change
        unBigMapId id0 @== unBigMapId id1
        -- The dupped bigmap has been assigned a new ID
        unBigMapId dupId @/= unBigMapId id1
    ]
  where
    -- Duplicate an existing bigmap, save it in the storage, and discard the old bigmap.
    dupAndDiscardOld :: Contract () (BigMap MText Integer) ()
    dupAndDiscardOld = L.mkContractWith L.intactCompilationOptions $ L.mkContractCode $
      L.cdr
      # L.dup
      # L.dip L.drop
      # L.nil # L.pair

    -- Duplicate an existing bigmap, save both (the old and the new) bigmaps in the storage.
    dupAndKeepBoth :: Contract () (BigMap MText Integer, BigMap MText Integer) ()
    dupAndKeepBoth = L.mkContractWith L.intactCompilationOptions $ L.mkContractCode $
      L.cdr
      # L.car
      # L.stackType @'[ BigMap MText Integer ]
      # L.dup
      # L.swap
      # L.pair
      # L.nil # L.pair

test_UniqueIDs :: TestTree
test_UniqueIDs =
  testScenario "all big_maps have unique IDs" $ scenario do
    let bm = [(1, 1)]
    let bigMaps =
          BigMaps
            bm
            (Just bm)
            [bm]
            (bm, bm)
            (Left bm)
            (Right bm)
            [(1, bm)]

    -- Originate the contract.
    -- Call the contract twice, to make sure the BigMapCounter is incremented
    -- inbetween operations.
    addr <- originate "create-big-maps" bigMaps createBigMaps
    transfer addr $ calling def bigMaps
    transfer addr $ calling def bigMaps

    finalStorage <- getStorage addr
    let bigMapIds = finalStorage ^.. biplate @_ @(BigMapId Natural Natural) . to unBigMapId

    -- The contract's storage should have 30 big_maps in total.
    --   * 8 created via the contract's initial storage,
    --   * For each contract call:
    --       * 8 that were passed in through the call's parameter
    --       * 1 via `EMPTY_BIG_MAP`, 1 via `DUP`, 1 via `DUP n`
    length bigMapIds @== 30

    assert (bigMapIds == ordNub bigMapIds) $
      unlinesF @[] @Doc
        [ "Expected all big_maps to have unique IDs, but some duplicates were found:"
        , blockListF (sort bigMapIds)
        ]

test_GetBigMapValue :: TestTree
test_GetBigMapValue =
  testScenario "getBigMapValue retrieves the correct value" $ scenario do
    contract1 <- originate
      "contract1"
      ([(1, "a"), (2, "b")], [(1, "c"), (2, "d")])
      $ idContract @() @(BigMap Integer MText, BigMap Integer MText)
    contract2 <- originate
      "contract1"
      ([(1, "e"), (2, "f")], [(1, "g"), (2, "h")])
      $ idContract @() @(BigMap Integer MText, BigMap Integer MText)

    (bigMap1, bigMap2) <- getStorage contract1
    (bigMap3, bigMap4) <- getStorage contract2

    getBigMapValue bigMap1 1 @@== "a"
    getBigMapValue bigMap1 2 @@== "b"
    getBigMapValue bigMap2 1 @@== "c"
    getBigMapValue bigMap2 2 @@== "d"
    getBigMapValue bigMap3 1 @@== "e"
    getBigMapValue bigMap3 2 @@== "f"
    getBigMapValue bigMap4 1 @@== "g"
    getBigMapValue bigMap4 2 @@== "h"

test_getBigMapSize :: TestTree
test_getBigMapSize =
  testScenario "getBigMapSize retrieves the correct size" $ scenario do
    let bmList = sampleList ++ [(5, "d")]
    c <- originate "contract" (mkBigMap bmList) $ idContract @() @(BigMap Integer MText)

    bigMapId <- getStorage c
    bmSize <- getBigMapSize bigMapId

    bmSize @== length bmList

test_getAllBigMapValues :: TestTree
test_getAllBigMapValues =
  testScenario "getAllBigMapValues retrieves the correct values" $ scenario do
    let bmValues = snd <$> sampleList
    c <- originate "contract" (mkBigMap sampleList) $ idContract @() @(BigMap Integer MText)

    bigMapId <- getStorage c
    vs <- getAllBigMapValues bigMapId

    sort vs @== sort bmValues

test_getAllBigMapValues_ConsistentWith_getBigMapValue :: TestTree
test_getAllBigMapValues_ConsistentWith_getBigMapValue =
  testScenario "results of 'getAllBigMapValues' and 'getBigMapValue' are consistent" $ scenario do
    c <- originate
      "contract"
      (mkBigMap sampleList)
      $ idContract @() @(BigMap Integer MText)

    bigMapId <- getStorage c

    vs <- forM sampleList $ \(k, _) -> getBigMapValue bigMapId k
    vs' <- getAllBigMapValues bigMapId

    sort vs @== sort vs'
    length vs @== length vs'

test_GetBigMapValueMaybe_InvalidBigMapId :: TestTree
test_GetBigMapValueMaybe_InvalidBigMapId =
  testScenario "getBigMapValueMaybe returns 'Nothing' when big_map ID is invalid" $ scenario do
    let invalidBigMapId = 2 ^ (99 :: Natural)
    getBigMapValueMaybe @Integer @Integer invalidBigMapId 0 @@== Nothing

test_GetBigMapValueMaybe_InvalidKey :: TestTree
test_GetBigMapValueMaybe_InvalidKey =
  testScenario "getBigMapValueMaybe returns 'Nothing' when key is invalid" $ scenario do
    addr <- originate
      "contract1"
      ([(1, "a"), (2, "b")], [(2, "c"), (3, "d")])
      $ idContract @() @(BigMap Integer MText, BigMap Integer MText)

    (bigMap1, bigMap2) <- getStorage addr

    getBigMapValueMaybe bigMap1 3 @@== Nothing
    getBigMapValueMaybe bigMap2 1 @@== Nothing

test_ID_BecomesInvalid :: TestTree
test_ID_BecomesInvalid =
  testScenario "big_map IDs become invalid once the big_map is removed from the storage" $ scenario do
    addr <- originate
      "contract1"
      [(1, "a")]
      $ saveInStorageContract @(BigMap Integer MText)

    bigMapId <- getStorage addr
    getBigMapValueMaybe bigMapId 1 @@== Just "a"

    -- Replace the big_map in the contract's storage with a new big_map (with same contents).
    transfer addr $ calling def [(1, "a")]

    -- The old big_map ID should now be invalid.
    getBigMapValueMaybe bigMapId 1 @@== Nothing

    -- The new big_map in the contract's storage should have a different ID than the first big_map.
    fmap unBigMapId (getStorage addr) @@/= unBigMapId bigMapId

sampleList :: [(Integer, MText)]
sampleList = [(1, "a"), (2, "b"), (3, "c"), (4, "d")]