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

module Test.Macro
  ( unit_PAPAIR
  , unit_UNPAPAIR
  , unit_CADR
  , unit_SET_CADR
  , unit_MAP_CADR
  , unit_mapPairLeaves
  , unit_expand
  , unit_expandValue
  , test_carnAndCdrnExpandToGetN
  ) where

import Hedgehog (forAll, property, (===))
import Hedgehog.Gen qualified as Gen
import Hedgehog.Gen.Michelson.Untyped qualified as Gen.U
import Hedgehog.Range qualified as Range
import Test.Hspec (Expectation, shouldBe)
import Test.Tasty (TestTree)
import Test.Tasty.Hedgehog (testProperty)

import Morley.Michelson.ErrorPos (InstrCallStack(..), LetName(..), SrcPos, srcPos)
import Morley.Michelson.Macro
import Morley.Michelson.Untyped
  (ExpandedOp(..), InstrAbstract(..), Value, Value'(..), mkAnnotation, noAnn)

import Test.Cleveland.Instances ()

defPos :: SrcPos
defPos = srcPos 1 1

defICS :: InstrCallStack
defICS = InstrCallStack [] defPos

-- TODO: it seems to me that these duplicated "where" blocks should be
-- replaced with some reasonable mini-EDSL - at least to facilitate tests
-- writing - and that would be a rather big refactoring.
-- Dunno how to deal with this duplication otherwise.
{-# ANN module ("HLint: ignore Reduce duplication" :: Text) #-}

unit_PAPAIR :: Expectation
unit_PAPAIR = do
  expandPapair defICS pair n n `shouldBe` [primEx $ PAIR n n n n]
  expandPapair defICS (P leaf pair) n n `shouldBe`
    [primEx $ DIP (expandMacro defICS $ PAPAIR pair n n), primEx $ PAIR n n n n]
  expandList [mac $ PAPAIR (P pair leaf) n n] `shouldBe`
    [WithSrcEx defICS $ SeqEx [primEx $ PAIR n n n n, primEx $ PAIR n n n n]]
  expandList [mac $ PAPAIR (P pair pair) n n] `shouldBe`
    [WithSrcEx defICS $ SeqEx [primEx (PAIR n n n n),
             primEx (DIP [primEx (PAIR n n n n)]),
             primEx (PAIR n n n n)]]
  where
    mac = flip Mac defPos
    primEx = PrimEx
    n = noAnn
    leaf = F n
    pair = P leaf leaf

unit_UNPAPAIR :: Expectation
unit_UNPAPAIR = do
  expandUnpapair defICS pair `shouldBe`
    [primEx $ UNPAIR n n n n]

  expandList [mac $ UNPAPAIR $ UP UF pair] `shouldBe`
    [WithSrcEx defICS $ SeqEx
      [ primEx (UNPAIR n n n n)
      , primEx (DIP [primEx (UNPAIR n n n n)])
      ]
    ]

  expandList [mac $ UNPAPAIR $ UP pair UF] `shouldBe`
    [WithSrcEx defICS $ SeqEx
      [ primEx (UNPAIR n n n n)
      , primEx (UNPAIR n n n n)
      ]
    ]

  expandList [mac $ UNPAPAIR $ UP pair pair] `shouldBe`
    [WithSrcEx defICS $ SeqEx
      [ primEx (UNPAIR n n n n)
      , primEx (DIP [primEx $ UNPAIR n n n n])
      , primEx (UNPAIR n n n n)
      ]
    ]
  where
    mac = flip Mac defPos
    primEx = PrimEx
    n = noAnn
    pair = UP UF UF

unit_CADR :: Expectation
unit_CADR = do
  expandCadr defICS ([A]) v f `shouldBe` [primEx $ CAR v f]
  expandCadr defICS ([D]) v f `shouldBe` [primEx $ CDR v f]
  expandCadr defICS (A:xs) v f `shouldBe` primEx (CAR n n) : expandMacro defICS (CADR xs v f)
  expandCadr defICS (D:xs) v f `shouldBe` primEx (CDR n n) : expandMacro defICS (CADR xs v f)
  where
    primEx = PrimEx
    v = "var"
    f = "field"
    n = noAnn
    xs = [A, D]

unit_SET_CADR :: Expectation
unit_SET_CADR = do
  expandSetCadr defICS [A] v f `shouldBe` primEx <$> [ DUP noAnn, CAR noAnn f, DROP
                                            , CDR "%%" noAnn, SWAP, PAIR noAnn v f "@"]
  expandSetCadr defICS [D] v f `shouldBe` primEx <$> [ DUP noAnn, CDR noAnn f, DROP
                                            , CAR "%%" noAnn, PAIR noAnn v "@" f]
  expandSetCadr defICS (A:xs) v f `shouldBe`
    primEx <$> [DUP noAnn, DIP (primEx carN : expandMacro defICS (SET_CADR xs noAnn f)), cdrN, SWAP, pairN]
  expandSetCadr defICS (D:xs) v f `shouldBe`
    primEx <$> [DUP noAnn, DIP (primEx cdrN : expandMacro defICS (SET_CADR xs noAnn f)), carN, pairN]
  where
    primEx = PrimEx
    v = "var"
    f = "field"
    xs = [A, D]
    carN = CAR "%%" noAnn
    cdrN = CDR "%%" noAnn
    pairN = PAIR noAnn v "@" "@"

unit_MAP_CADR :: Expectation
unit_MAP_CADR = do
  expandMapCadr defICS [A] v f ops `shouldBe`
    primEx <$> [DUP noAnn, cdrN, DIP [primEx $ CAR noAnn f, SeqEx ops'], SWAP, PAIR noAnn v f "@"]
  expandMapCadr defICS [D] v f ops `shouldBe`
    concat [primEx <$> [DUP noAnn, CDR noAnn f], [SeqEx ops'], primEx <$> [SWAP, carN, PAIR noAnn v "@" f]]
  expandMapCadr defICS (A:xs) v f ops `shouldBe`
    primEx <$> [DUP noAnn, DIP (primEx carN : expandMacro defICS (MAP_CADR xs noAnn f ops)), cdrN, SWAP, pairN]
  expandMapCadr defICS (D:xs) v f ops `shouldBe`
    primEx <$> [DUP noAnn, DIP (primEx cdrN : expandMacro defICS (MAP_CADR xs noAnn f ops)), carN, pairN]
  where
    primEx = PrimEx
    v = "var"
    f = "field"
    n = noAnn
    xs = [A, D]
    ops = [Prim (DUP n) defPos]
    ops' = [WithSrcEx defICS $ PrimEx (DUP n)]
    carN = CAR "%%" noAnn
    cdrN = CDR "%%" noAnn
    pairN = PAIR noAnn v "@" "@"

unit_mapPairLeaves :: Expectation
unit_mapPairLeaves = do
  mapPairLeaves [f, f] pair `shouldBe` P (F f) (F f)
  mapPairLeaves annotations (P pair (F n)) `shouldBe`
    P (P (leaf "field1") (leaf "field2")) (leaf "field3")
  mapPairLeaves annotations (P pair pair) `shouldBe`
    P (P (leaf "field1") (leaf "field2")) (P (leaf "field3") (F n))
  where
    annotations = unsafe . mkAnnotation <$> ["field1", "field2", "field3"]
    n = noAnn
    f = "field"
    leaf f' = F (unsafe . mkAnnotation $ f')
    pair = P (F n) (F n)

unit_expand :: Expectation
unit_expand = do
  expand [LetName "a"] diip `shouldBe` expandedDiip
  expand [LetName "a"] (prim $ IF [diip] [diip]) `shouldBe` (primEx $ IF [expandedDiip] [expandedDiip])
  expand [LetName "a"] (Seq [diip, diip] defPos) `shouldBe` (WithSrcEx aIcs $ SeqEx $ [expandedDiip, expandedDiip])
  where
    aIcs = InstrCallStack [LetName "a"] defPos
    prim = flip Prim defPos
    primEx = WithSrcEx aIcs . PrimEx
    mac = flip Mac defPos
    diip :: ParsedOp
    diip = mac (DIIP 2 [prim SWAP])
    expandedDiip :: ExpandedOp
    expandedDiip = primEx (DIPN 2 [primEx SWAP])

unit_expandValue :: Expectation
unit_expandValue = do
  expandValue parsedPair `shouldBe` expandedPair
  expandValue parsedPapair `shouldBe` expandedPapair
  expandValue parsedLambdaWithMac `shouldBe` expandedLambdaWithMac
  where
    mac = flip Mac defPos
    primEx = PrimEx

    parsedPair :: Value' ParsedOp
    parsedPair = ValuePair (ValueInt 5) (ValueInt 5)

    expandedPair :: Value
    expandedPair = ValuePair (ValueInt 5) (ValueInt 5)

    parsedPapair :: Value' ParsedOp
    parsedPapair = ValuePair (ValuePair (ValueInt 5) (ValueInt 5)) (ValueInt 5)

    expandedPapair :: Value
    expandedPapair = ValuePair (ValuePair (ValueInt 5) (ValueInt 5)) (ValueInt 5)

    parsedLambdaWithMac :: Value' ParsedOp
    parsedLambdaWithMac = ValueLambda $
      one (mac (PAPAIR (P (F noAnn) (P (F noAnn) (F noAnn))) noAnn noAnn))

    expandedLambdaWithMac :: Value
    expandedLambdaWithMac = ValueLambda . one $ WithSrcEx defICS $ SeqEx
      [ primEx $ DIP [primEx $ PAIR noAnn noAnn noAnn noAnn]
      , primEx $ PAIR noAnn noAnn noAnn noAnn
      ]

test_carnAndCdrnExpandToGetN :: [TestTree]
test_carnAndCdrnExpandToGetN =
  [ testProperty "CAR k to GET 2k+1" $ property do
    k <- forAll $ Gen.word $ Range.linear 0 100
    n <- forAll $ Gen.U.genAnnotation
    expand' (CARN n k) === primEx (GETN n $ 2 * k + 1)
  , testProperty "CDR k to GET 2k" $ property do
    k <- forAll $ Gen.word $ Range.linear 0 100
    n <- forAll $ Gen.U.genAnnotation
    expand' (CDRN n k) === primEx (GETN n $ 2 * k)
  ]
  where
    primEx = WithSrcEx defICS . SeqEx . one . PrimEx
    mac = flip Mac defPos
    expand' = expand [] . mac