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

module Test.Ext
  ( test_PRINT_and_TEST_ASSERT
  , test_STACKTYPE
  ) where

import Data.Default (def)
import Fmt (pretty)
import Test.HUnit (Assertion, assertFailure)
import Test.Hspec.Expectations (shouldSatisfy)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (testCase)

import Morley.Michelson.Interpret (MorleyLogs(..), interpret)
import Morley.Michelson.Runtime.Dummy (dummyBigMapCounter, dummyContractEnv, dummyGlobalCounter)
import Morley.Michelson.TypeCheck
  (HST(..), SomeHST(..), TypeCheckMode(..), mkSomeParamType, runTypeCheck, throwingTCError,
  typeCheckExt, typeCheckInstr, typeCheckingWith, unsafeWithWTP)
import Morley.Michelson.Typed (epcPrimitive, pattern (:#), pattern AsUType)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Untyped
  (ExpandedExtInstr, ExtInstrAbstract(..), ParameterType(..), StackTypePattern(..), T(..), Ty(..),
  TyVar(..), noAnn)
import Morley.Util.PeanoNatural (PeanoNatural(..))
import Test.Cleveland.Instances ()

testassert_square :: T.Contract ('T.TPair 'T.TInt 'T.TInt) 'T.TUnit
testassert_square = T.defaultContract $
  T.CAR :#
  T.UNPAIR :#
  T.Ext (T.TEST_ASSERT
    (T.TestAssert "CheckSides"
      (T.PrintComment [Left "Sides are ", Right (T.StackRef Zero), Left " x ", Right (T.StackRef One)])
      ( T.DUP :#
        T.GT :#
        T.DIP (T.PUSH (T.VInt 101) :# T.COMPARE :# T.GT) :#
        T.DIPN Two (T.DUP :# T.GT :# T.DIP (T.PUSH (T.VInt 101) :# T.COMPARE :# T.GT)) :#
        T.AND :# T.AND :# T.AND
      )
    )) :#
  T.MUL :#
  T.Ext (T.PRINT $ T.PrintComment [Left "Area is ", Right $ T.StackRef Zero]) :#
  T.DROP :#
  T.UNIT :#
  T.NIL @'T.TOperation :#
  T.PAIR

-- | This is like 'testassert_square', but test_assert is executed on stack
-- with more than one element.
testassert_square2 :: T.Contract ('T.TPair 'T.TInt 'T.TInt) 'T.TUnit
testassert_square2 = T.defaultContract $
  T.CAR :#
  T.DIP T.UNIT :#
  T.UNPAIR :#
  T.Ext (T.TEST_ASSERT
    (T.TestAssert "CheckSides"
      (T.PrintComment [Left "Sides are ", Right (T.StackRef Zero), Left " x ", Right (T.StackRef One)])
      ( T.DUP :#
        T.GT :#
        T.DIP (T.PUSH (T.VInt 101) :# T.COMPARE :# T.GT) :#
        T.DIPN Two (T.DUP :# T.GT :# T.DIP (T.PUSH (T.VInt 101) :# T.COMPARE :# T.GT)) :#
        T.AND :# T.AND :# T.AND
      )
    )) :#
  T.MUL :#
  T.Ext (T.PRINT $ T.PrintComment [Left "Area is ", Right $ T.StackRef Zero]) :#
  T.DROP :#
  T.NIL @'T.TOperation :#
  T.PAIR

test_PRINT_and_TEST_ASSERT :: [TestTree]
test_PRINT_and_TEST_ASSERT =
  [ testGroup "testassert_square" $ testAssertSquareImpl testassert_square
  , testGroup "testassert_square2" $ testAssertSquareImpl testassert_square2
  ]
  where
    testAssertSquareImpl c =
      [ testCase "TEST_ASSERT assertion passed" $ do
        runTest True c 100 100
        runTest True c 1 1
      , testCase "TEST_ASSERT assertion failed" $ do
        runTest False c 0 100
        runTest False c -1 -2
      ]

    runTest corr contract x y = do
      let x' = T.VInt x :: T.Value 'T.TInt
      let y' = T.VInt y :: T.Value 'T.TInt
      let area' = T.VInt $ x * y :: T.Value 'T.TInt
      let check (a, (_, s)) =
            if corr then isRight a && s == MorleyLogs ["Area is " <> pretty area']
            else isLeft a && s == MorleyLogs ["Sides are " <> pretty x' <> " x " <> pretty y']
      interpret
        contract
        epcPrimitive
        (T.VPair (x', y'))
        T.VUnit
        dummyGlobalCounter
        dummyBigMapCounter
        dummyContractEnv
          `shouldSatisfy` check

test_STACKTYPE :: [TestTree]
test_STACKTYPE =
  [ testCase "Correct test on [] pattern" $ runExtTest test1 True
  , testCase "Correct test on [a, b] pattern" $ runExtTest test2 True
  , testCase "Correct test on [a, b, ...] pattern" $ runExtTest test3 True
  , testCase "Correct test on [a, b, ...] pattern and stack [a, b]" $ runExtTest test4 True

  , testCase "Failed test on [] pattern and stack [a]" $ runExtTest test5 False
  , testCase "Failed test on [a, b] pattern and stack [a, b, c]" $ runExtTest test6 False
  , testCase "Failed test on [a, b] pattern and stack [a]" $ runExtTest test7 False
  , testCase "Failed test on [a, b, ...] pattern and stack [a]" $ runExtTest test8 False
  , testCase "Failed test on [a, b] pattern and stack [a, c]" $ runExtTest test9 False
  ]
  where
    p2 = StkCons (TyCon t1) (StkCons (TyCon t2) StkEmpty)
    p3 = StkCons (TyCon t1) (StkCons (TyCon t2) StkRest)

    test1 = (STACKTYPE StkEmpty, convertToHST [])
    test2 = (STACKTYPE p2, convertToHST [t1, t2])
    test3 = (STACKTYPE p3, convertToHST [t1, t2, t3])
    test4 = (STACKTYPE p3, convertToHST [t1, t2])

    test5 = (STACKTYPE StkEmpty, convertToHST [t1])
    test6 = (STACKTYPE p2, convertToHST [t1, t2, t3])
    test7 = (STACKTYPE p2, convertToHST [t1])
    test8 = (STACKTYPE p3, convertToHST [t1])
    test9 = (STACKTYPE p2, convertToHST [t1, t3])

    t1 = Ty (TOption (Ty TKey "key")) "opt"
    t2 = Ty (TPair "f" "s" noAnn noAnn (Ty TUnit "x") (Ty TSignature "s")) noAnn
    t3 = Ty TInt "tint"

    convertToHST :: [Ty] -> SomeHST
    convertToHST [] = SomeHST SNil
    convertToHST (AsUType (nt :: T.Notes t1) : ts) =
      case convertToHST ts of
        SomeHST is -> unsafeWithWTP @t1 $ SomeHST ((nt, T.Dict, noAnn) ::& is)

    nh (ni, si) =
      typeCheckingWith def $
      runTypeCheck (TypeCheckContract $
                      unsafe . mkSomeParamType $
                      ParameterType (Ty TKey noAnn) noAnn)
                   (usingReaderT def $
                      throwingTCError $
                      typeCheckExt typeCheckInstr ni si)

    runExtTest :: (ExpandedExtInstr, SomeHST) -> Bool -> Assertion
    runExtTest (ui, SomeHST hst) correct = case (nh (ui, hst), correct) of
      (Right _, False) -> assertFailure $ "Test expected to fail but it passed"
      (Left e, True)   -> assertFailure $ "Test expected to pass but it failed with error: " <> pretty e
      _                -> pass