{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE QuasiQuotes         #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Defines tests of STG programs that are based on marshalling a value into
-- the STG, forcing a value, and marshalling that value out again for comparison
-- with a reference.
module Test.Machine.Evaluate.TestTemplates.MarshalledValue (
    MarshalledValueTestSpec(..),
    MarshalSourceSpec(..),
    defSpec,
    marshalledValueTest,
) where



import           Data.List.NonEmpty                    (NonEmpty (..))
import           Data.Text                             (Text)
import qualified Data.Text                             as T
import           Data.Text.Prettyprint.Doc
import           Data.Text.Prettyprint.Doc.Render.Text

import Stg.Language
import Stg.Language.Prettyprint
import Stg.Machine
import Stg.Machine.Types
import Stg.Marshal
import Stg.Parser.QuasiQuoter   (stg)

import Test.Orphans          ()
import Test.Tasty
import Test.Tasty.QuickCheck



-- | Specifies a test that is based on marshalling a value out of the STG and
-- comparing it to a known value.
data MarshalledValueTestSpec input output = MarshalledValueTestSpec
    { testName :: Text
        -- ^ The reference function's name. Used only for display purposes.

    , failPredicate :: StgState -> Bool
        -- ^ Fail if this predicate holds. This can be used to constrain the
        -- heap size during the test, for example.

    , sourceSpec :: input -> MarshalSourceSpec output
        --  * STG program to run

    , maxSteps :: Integer
        -- ^ Maximum number of steps to take

    , failWithInfo :: Bool
        -- ^ Print program code and final state on test failure?
    }

data MarshalSourceSpec output = MarshalSourceSpec
    { resultVar     :: Var      -- ^ value to observe the value of, e.g. @main@
    , expectedValue :: output   -- ^ expected result value
    , source        :: Program  -- ^ STG program to run
    }

defSpec :: MarshalledValueTestSpec input b
defSpec = MarshalledValueTestSpec
    { testName = "Default Haskell reference test spec template"
    , maxSteps = 1024
    , failWithInfo = False
    , failPredicate = const False
    , sourceSpec = \_ -> MarshalSourceSpec
        { resultVar = "main"
        , expectedValue = error "No expected value generator set in test"
        , source = [stg| main = \ -> DummySource |] }}

marshalledValueTest
    :: forall input output.
       ( Show input, Arbitrary input
       , Eq output, Show output, FromStg output, PrettyStgi output )
    => MarshalledValueTestSpec input output
    -> TestTree
marshalledValueTest testSpec = testProperty (T.unpack (testName testSpec)) test
  where
    test :: ( Show input, Arbitrary input
            , Eq output, Show output, FromStg output, PrettyStgi output )
         => input
         -> Property
    test input =
        let program = initialState "main" (source (sourceSpec testSpec input))
            states = evalsUntil
                (RunForMaxSteps (maxSteps testSpec))
                (HaltIf (const False))
                (PerformGc (const Nothing))
                program
            verifyLoop (state :| _)
                | failPredicate testSpec state =
                    fail_failPredicateTrue testSpec input state
            verifyLoop (state :| rest) = case fromStg state (resultVar (sourceSpec testSpec input)) of
                Left err -> case err of
                    TypeMismatch -> fail_typeMismatch testSpec input state
                    IsBlackhole -> continue state rest
                    IsWrongLambdaType LambdaFun -> fail_functionValue testSpec input state
                    IsWrongLambdaType LambdaThunk -> continue state rest
                    IsWrongLambdaType LambdaCon -> error
                        "Critial error in test: found a constructor, expected\
                        \ a constructor, but still ran into the failure case\
                        \ somehow. Please report this as a bug."
                    BadArity -> fail_conArity testSpec input state
                    NotFound{} -> fail_notFound testSpec input state
                    AddrNotOnHeap -> fail_addrNotOnHeap testSpec input state
                    NoConstructorMatch -> fail_NoConstructorMatch testSpec input state
                Right actualValue -> assertEqual actualValue testSpec input state
            continue lastState = \case
                [] -> fail_valueNotFound testSpec input lastState
                (x:xs) -> verifyLoop (x :| xs)

        in verifyLoop states

assertEqual
    :: forall input output. (Eq output, PrettyStgi output)
    => output
    -> MarshalledValueTestSpec input output
    -> input
    -> StgState
    -> Property
assertEqual
    actual
    testSpec
    input
    finalState
  = counterexample (failText expected) (actual == expected)
  where
    expected :: output
    expected = expectedValue (sourceSpec testSpec input)

    failText :: PrettyStgi a => a -> String
    failText ex = (T.unpack . renderStrict . layoutPretty defaultLayoutOptions . vsep)
        [ "Machine produced an invalid result."
        , "Expected:" <+> (prettyStgi ex :: Doc StgiAnn)
        , "Actual:  " <+> (prettyStgi actual :: Doc StgiAnn)
        , if failWithInfo testSpec
            then vsep
                [ hang 4 (vsep ["Program:", prettyStgi (source (sourceSpec testSpec input))])
                , hang 4 (vsep ["Final state:", prettyStgi finalState]) ]
            else failWithInfoInfoText ]

failWithInfoInfoText :: Doc ann
failWithInfoInfoText = "Run test case with failWithInfo to see the final state."

fail_template
    :: Doc StgiAnn
    -> MarshalledValueTestSpec input a
    -> input
    -> StgState
    -> Property
fail_template
    failMessage
    testSpec
    input
    finalState
  = counterexample failText False
  where
    failText = (T.unpack . renderStrict . layoutPretty defaultLayoutOptions . vsep)
        [ failMessage
        , "Final machine state info:"
            <+> (prettyStgi (stgInfo finalState) :: Doc StgiAnn)
        , if failWithInfo testSpec
            then vsep
                [ hang 4 (vsep ["Program:", prettyStgi (source (sourceSpec testSpec input)) :: Doc StgiAnn])
                , hang 4 (vsep ["Final state:", prettyStgi finalState :: Doc StgiAnn]) ]
            else failWithInfoInfoText ]

fail_failPredicateTrue, fail_valueNotFound, fail_typeMismatch, fail_conArity,
    fail_notFound, fail_addrNotOnHeap, fail_NoConstructorMatch, fail_functionValue
    :: MarshalledValueTestSpec a b -> a -> StgState -> Property
fail_failPredicateTrue  = fail_template "Failure predicate held for an intemediate state"
fail_valueNotFound      = fail_template "None of the machine states produced a (marshallable)\
                                        \ value to compare the expected value to"
fail_typeMismatch       = fail_template "Type mismatch in input/expected output"
fail_conArity           = fail_template "Bad constructor arity in created value"
fail_notFound           = fail_template "Variable not found"
fail_addrNotOnHeap      = fail_template "Address not found on heap"
fail_NoConstructorMatch = fail_template "No constructor match"
fail_functionValue      = fail_template "Function value encountered; can only do algebraic"