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

module Test.Michelson.Typed.AnnotatedValue
  ( unit_variant_read
  , unit_variant_update
  , unit_field_read
  , unit_field_update
  , unit_nodes
  , unit_asList
  , unit_asMap
  )
  where

import Control.Lens (_Just, each, ix, (+~), (<>~), (^@..))
import Fmt (Buildable, GenericBuildable(..))
import Test.Cleveland.Instances ()
import Test.Tasty.HUnit ((@?=))

import Morley.Michelson.Parser (notes)
import Morley.Michelson.Text (MText)
import Morley.Michelson.Typed.AnnotatedValue
import Morley.Michelson.Typed.Annotation (Notes(..), starNotes)
import Morley.Michelson.Typed.Haskell.Value (BigMap, IsoValue, ToT, toVal)

----------------------------------------------------------------------------
-- Sum types
----------------------------------------------------------------------------

data Param
  = P1 Integer
  | P2 MText
  | P3 (MText, Natural)
  | P4 NestedParam
  deriving stock Generic
  deriving anyclass (IsoValue)
  deriving Buildable via (GenericBuildable Param)

data NestedParam
  = NP1 Integer
  | NP2 Natural
  deriving stock Generic
  deriving anyclass (IsoValue)
  deriving Buildable via (GenericBuildable NestedParam)

paramNotes :: Notes (ToT Param)
paramNotes =
  [notes|
    or
      (or (int %p1)
          (string %p2)
      )
      (or (pair %p3 string nat)
          (or %p4 (int %p1) (nat %p2))
      )
  |]

nestedParamNotes :: Notes (ToT NestedParam)
nestedParamNotes = [notes|or (int %p1) (nat %p2)|]

unit_variant_read :: IO ()
unit_variant_read = do
  let mkParam param = SomeAnnotatedValue paramNotes (toVal param)

  let p1 = mkParam $ P1 1
  p1 ^.. variant "p1" @?= one (SomeAnnotatedValue starNotes (toVal @Integer 1))
  p1 ^.. variant "p2" @?= []
  p1 ^.. variant "p3" @?= []
  p1 ^.. variant "p4" @?= []

  let p2 = mkParam $ P2 "a"
  p2 ^.. variant "p1" @?= []
  p2 ^.. variant "p2" @?= one (SomeAnnotatedValue starNotes (toVal @MText "a"))
  p2 ^.. variant "p3" @?= []
  p2 ^.. variant "p4" @?= []

  let p3 = mkParam $ P3 ("a", 1)
  p3 ^.. variant "p1" @?= []
  p3 ^.. variant "p2" @?= []
  p3 ^.. variant "p3" @?= one (SomeAnnotatedValue starNotes (toVal @(MText, Natural) ("a", 1)))
  p3 ^.. variant "p4" @?= []

  let p4_1 = mkParam $ P4 (NP1 1)
  p4_1 ^.. variant "p1" @?= []
  p4_1 ^.. variant "p2" @?= []
  p4_1 ^.. variant "p3" @?= []
  p4_1 ^.. variant "p4" @?= one (SomeAnnotatedValue nestedParamNotes (toVal (NP1 1)))
  p4_1 ^.. variant "p4" . variant "p1" @?= one (SomeAnnotatedValue starNotes (toVal @Integer 1))
  p4_1 ^.. variant "p4" . variant "p2" @?= []

  let p4_2 = mkParam $ P4 (NP2 1)
  p4_2 ^.. variant "p1" @?= []
  p4_2 ^.. variant "p2" @?= []
  p4_2 ^.. variant "p3" @?= []
  p4_2 ^.. variant "p4" @?= one (SomeAnnotatedValue nestedParamNotes (toVal (NP2 1)))
  p4_2 ^.. variant "p4" . variant "p1" @?= []
  p4_2 ^.. variant "p4" . variant "p2" @?= one (SomeAnnotatedValue starNotes (toVal @Natural 1))

unit_variant_update :: IO ()
unit_variant_update = do
  let mkParam param = SomeAnnotatedValue paramNotes (toVal param)

  (mkParam (P1 1) & variant "p1" . castTo @Integer . value +~ 99)
    @?= mkParam (P1 100)

  (mkParam (P2 "a") & variant "p2" . castTo @MText . value <>~ "!")
    @?= mkParam (P2 "a!")

  (mkParam (P3 ("a", 1)) & variant "p3" . castTo @(MText, Natural) . value . _1 <>~ "!")
    @?= mkParam (P3 ("a!", 1))

  (mkParam (P4 (NP1 1)) & variant "p4" . variant "p1" . castTo @Integer . value +~ 99)
    @?= mkParam (P4 (NP1 100))

  (mkParam (P4 (NP2 1)) & variant "p4" . variant "p2" . castTo @Natural . value +~ 99)
    @?= mkParam (P4 (NP2 100))

----------------------------------------------------------------------------
-- Product types
----------------------------------------------------------------------------

data Storage = Storage
  { sField1 :: Integer
  , sField2 :: (MText, Natural)
  , sField3 :: Maybe Integer
  , sField4 :: Nested
  }
  deriving stock Generic
  deriving anyclass (IsoValue)
  deriving Buildable via (GenericBuildable Storage)

data Nested = Nested
  { nField1 :: MText
  , nField2 :: Integer
  }
  deriving stock Generic
  deriving anyclass (IsoValue)
  deriving Buildable via (GenericBuildable Nested)

storageNotes :: Notes (ToT Storage)
storageNotes =
  [notes|
    pair
      (pair
        (int %field1)
        (pair %field2 string nat)
      )
      (pair
        (option %field3 int)
        (pair %field4
          (string %field1)
          (int %field2)
        )
      )
  |]

nestedNotes :: Notes (ToT Nested)
nestedNotes = [notes|pair (string %field1) (int %field2)|]

storage :: Storage
storage = Storage
  { sField1 = 1
  , sField2 = ("a", 2)
  , sField3 = Just 3
  , sField4 = Nested
    { nField1 = "b"
    , nField2 = 4
    }
  }

storageSav :: SomeAnnotatedValue
storageSav = SomeAnnotatedValue storageNotes (toVal storage)

unit_field_read :: IO ()
unit_field_read = do
  storageSav ^@.. fields @?=
    [ ("field1", SomeAnnotatedValue starNotes (toVal $ sField1 storage))
    , ("field2", SomeAnnotatedValue starNotes (toVal $ sField2 storage))
    , ("field3", SomeAnnotatedValue starNotes (toVal $ sField3 storage))
    , ("field4", SomeAnnotatedValue nestedNotes (toVal $ sField4 storage))
    ]

  storageSav ^@.. field "field4" . fields @?=
    [ ("field1", SomeAnnotatedValue starNotes (toVal $ nField1 $ sField4 storage))
    , ("field2", SomeAnnotatedValue starNotes (toVal $ nField2 $ sField4 storage))
    ]

unit_field_update :: IO ()
unit_field_update = do
  (storageSav & field "field1" . castTo @Integer . value +~ 100)
    @?= SomeAnnotatedValue storageNotes
        (toVal $ storage { sField1 = 101 })

  (storageSav & field "field2" . castTo @(MText, Natural) . value . _2 +~ 100)
    @?= SomeAnnotatedValue storageNotes
        (toVal $ storage { sField2 = ("a", 102) })

  (storageSav & field "field3" . castTo @(Maybe Integer) . value . _Just +~ 100)
    @?= SomeAnnotatedValue storageNotes
        (toVal $ storage { sField3 = Just 103 })

  (storageSav & field "field4" . field "field1" . castTo @MText . value <>~ "!")
    @?= SomeAnnotatedValue storageNotes
        (toVal $ storage { sField4 = (sField4 storage) { nField1 = "b!" } })

  (storageSav & field "field4" . field "field2" . castTo @Integer . value +~ 100)
    @?= SomeAnnotatedValue storageNotes
        (toVal $ storage { sField4 = (sField4 storage) { nField2 = 104 } })

unit_nodes :: IO ()
unit_nodes = do
  storageSav ^@.. nodes @?=
    [ ("", SomeAnnotatedValue
          [notes|pair (int %field1) (pair %field2 string nat)|]
          (toVal @(Integer, (MText, Natural)) $ (1, ("a", 2)))
      )
    , ("", SomeAnnotatedValue
          [notes|pair (option %field3 int) (pair %field4 (string %field1) (int %field2))|]
          (toVal @(Maybe Integer, (MText, Integer)) $ (Just 3, ("b", 4)))
      )
    , ("field1", SomeAnnotatedValue starNotes (toVal $ sField1 storage))
    , ("field2", SomeAnnotatedValue starNotes (toVal $ sField2 storage))
    , ("", SomeAnnotatedValue starNotes (toVal @MText "a"))
    , ("", SomeAnnotatedValue starNotes (toVal @Natural 2))
    , ("field3", SomeAnnotatedValue starNotes (toVal $ sField3 storage))
    , ("field4", SomeAnnotatedValue nestedNotes (toVal $ sField4 storage))
    , ("field1", SomeAnnotatedValue starNotes (toVal $ nField1 $ sField4 storage))
    , ("field2", SomeAnnotatedValue starNotes (toVal $ nField2 $ sField4 storage))
    ]

  storageSav ^.. node "field1" @?=
    [ SomeAnnotatedValue starNotes (toVal $ sField1 storage)
    , SomeAnnotatedValue starNotes (toVal $ nField1 $ sField4 storage)
    ]

----------------------------------------------------------------------------
-- Collections
----------------------------------------------------------------------------

data Collections = Collections
  { cListField :: [Integer]
  , cMapField :: Map MText Natural
  , cBigMapField :: BigMap MText Natural
  }
  deriving stock Generic
  deriving anyclass (IsoValue)
  deriving Buildable via (GenericBuildable Collections)

collectionsNotes :: Notes (ToT Collections)
collectionsNotes =
  [notes|
    pair
      (list %listField int)
      (map %mapField string nat)
      (big_map %bigMapField string nat)
  |]

collectionsSav :: SomeAnnotatedValue
collectionsSav =
  SomeAnnotatedValue
    collectionsNotes
    (toVal Collections
      { cListField = [1, 2, 3]
      , cMapField = one ("b", 3)
      , cBigMapField = one ("c", 4)
      }
    )

unit_asList :: IO ()
unit_asList = do
  collectionsSav ^.. field "listField" . asList . each . castTo @Integer . value
    @?= [1, 2, 3]

unit_asMap :: IO ()
unit_asMap = do
  collectionsSav ^? field "mapField" . asMap @MText . ix "b" . castTo @Natural . value
    @?= Just 3
  collectionsSav ^? field "bigMapField" . asMap @MText . ix "c" . castTo @Natural . value
    @?= Just 4