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

{-# LANGUAGE QuantifiedConstraints #-}

-- | Testing utility functions used by testing framework itself or
-- intended to be used by test writers.

module Test.Cleveland.Util
  ( leftToShowPanic
  , (?-)

  -- * Property
  , failedTest
  , succeededTest
  , eitherIsLeft
  , eitherIsRight
  , meanTimeUpperBoundProp
  , meanTimeUpperBoundPropNF

  -- * Generator
  , genTuple2
  , genRandom
  , runGen

  -- * Roundtrip
  , roundtripTree
  , assertGoesBefore
  , goesBefore

  -- * Pretty-printing
  , formatValue
  , formatSomeValue
  , ShowWith(..)
  , Showing(..)

  -- * Time
  , ceilingUnit
  , timeToFixed
  , timeToNominalDiffTime

  -- * Bytes
  , stripOptional0x
  , fromHex
  , parseAddressFromHex

  -- * Traversals
  , mapEach
  , forEach

  -- * Re-exports
  --
  -- | These functions from @Time@ are re-exported here to make it convenient to call
  -- 'meanTimeUpperBoundProp' and 'meanTimeUpperBoundPropNF'.
  , mcs, ms, sec, minute
  ) where

import Debug qualified (show)

import Control.Lens qualified as L
import Control.Monad.Random (MonadRandom, evalRand, mkStdGen)
import Criterion (Benchmarkable, benchmarkWith', nf, whnf)
import Criterion.Main (defaultConfig)
import Criterion.Types (SampleAnalysis(anMean), Verbosity(Quiet), reportAnalysis, verbosity)
import Data.Fixed (Fixed, HasResolution)
import Data.Ratio ((%))
import Data.Singletons (demote)
import Data.Text qualified as T
import Data.Time (NominalDiffTime, secondsToNominalDiffTime)
import Data.Typeable (typeRep)
import Fmt (Buildable, Builder, build, pretty, (+|), (|+))
import Hedgehog
  (Gen, MonadGen, MonadTest, Property, annotate, evalIO, failure, forAll, property, success,
  tripping, withTests)
import Hedgehog.Gen qualified as Gen
import Hedgehog.Internal.Gen (runGenT)
import Hedgehog.Internal.Seed qualified as Seed
import Hedgehog.Internal.Tree (TreeT(runTreeT), nodeValue)
import Hedgehog.Range qualified as Range
import Statistics.Types (Estimate(estPoint))
import Test.HUnit (Assertion, assertFailure)
import Test.Tasty (TestTree)
import Test.Tasty.HUnit (testCase)
import Test.Tasty.Hedgehog (testProperty)
import Text.Hex (decodeHex)
import Text.Printf (printf)
import Text.Show qualified
import Time
  (KnownDivRat, KnownUnitName, Microsecond, Millisecond, Minute, Nanosecond, Picosecond, Rat,
  RatioNat, Second, Time, mcs, minute, ms, ns, sec, time, timeout, toUnit, unTime, unitNameVal)

import Morley.Michelson.Doc (DocItem, docItemPosition)
import Morley.Michelson.Typed (SingI, SomeConstrainedValue)
import Morley.Michelson.Typed qualified as T
import Morley.Tezos.Address
import Morley.Util.Constrained

import Test.Cleveland.Instances ()

{-
Note [Hedgehog & withFrozenCallStack]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

When an Hedgehog test fails, Hedgehog displays the source code for all callstack frames,
and annotates the top-most frame with the error message.

When an assertion helper like `failedTest` fails, the helper itself will be the top-most frame.
This means that, under normal circumstances, Hedgehog would display the error message
next to `failedTest`'s source code, and **not** next to the user's source code.

In order to force Hedgehog to display the error message next to the user's source code,
we use `withFrozenCallStack` to effectively remove the assertion helper from the callstack.

On a side note: the reason why we don't need to use `withFrozenCallStack` in Cleveland's
assertion helpers is because Cleveland only displays the bottom-most callstack
frame (i.e., the user's source code).
See: Note [Cleveland & callstacks] for more information.
-}

leftToShowPanic :: (PrettyShow e, Show e, HasCallStack) => Either e a -> a
leftToShowPanic = either (error . show) id

-- | Make a tuple with name without extra syntactic noise.
(?-) :: Text -> a -> (Text, a)
(?-) = (,)
infixr 0 ?-

----------------------------------------------------------------------------
-- Property
----------------------------------------------------------------------------

-- | A 'Property' that always fails with given message.
failedTest :: (HasCallStack, MonadTest m) => Text -> m ()
failedTest r =
  -- See: Note [Hedgehog & withFrozenCallStack]
  withFrozenCallStack $ annotate (toString r) >> failure

-- | A 'Property' that always succeeds.
succeededTest :: MonadTest m => m ()
succeededTest = success

-- | The 'Property' holds on `Left a`.
eitherIsLeft :: (Show b, MonadTest m, HasCallStack) => Either a b -> m ()
eitherIsLeft = \case
  Left _ -> succeededTest
  Right x ->
    -- See: Note [Hedgehog & withFrozenCallStack]
    withFrozenCallStack $ failedTest $ "expected Left, got Right (" <> Debug.show x <> ")"

-- | The 'Property' holds on `Right b`.
eitherIsRight :: (Show a, MonadTest m, HasCallStack) => Either a b -> m ()
eitherIsRight = \case
  Right _ -> succeededTest
  Left x ->
    -- See: Note [Hedgehog & withFrozenCallStack]
    withFrozenCallStack $ failedTest $ "expected Right, got Left (" <> Debug.show x <> ")"

-- | Benchmarks the given function and checks that the mean time to evaluate to weak head
-- normal form is under the given amount of time.
--
-- This test fails if the benchmark takes longer than 30 seconds to run.
meanTimeUpperBoundProp
  :: (KnownDivRat unit Second, KnownUnitName unit, HasCallStack)
  => Time unit -> (a -> b) -> a -> Property
meanTimeUpperBoundProp upperBound run arg =
  -- See: Note [Hedgehog & withFrozenCallStack]
  withFrozenCallStack $
    checkReport upperBound $ whnf run arg

-- | Benchmarks the given function and checks that the mean time to evaluate to
-- normal form is under the given amount of time.
--
-- This test aborts and fails if the benchmark takes longer than 120 seconds to run.
meanTimeUpperBoundPropNF
  :: (KnownDivRat unit Second, KnownUnitName unit, HasCallStack, NFData b)
  => Time unit -> (a -> b) -> a -> Property
meanTimeUpperBoundPropNF upperBound run arg =
  -- See: Note [Hedgehog & withFrozenCallStack]
  withFrozenCallStack $
    checkReport upperBound $ nf run arg

checkReport
  :: (KnownDivRat unit Second, KnownUnitName unit)
  => HasCallStack => Time unit -> Benchmarkable -> Property
checkReport upperBound benchmarkable =
  withTests 1 $ property $
    evalIO runBench >>= \case
      Nothing -> failedTest "Expected benchmark to complete within 120 seconds."
      Just report ->
        let mean = sec . realToFrac @Double @RatioNat . estPoint . anMean $ reportAnalysis report
        in  if mean < toUnit @Second upperBound
              then succeededTest
              else failedTest $
                "Expected mean estimate to be under "
                <> show upperBound
                <> ", but was "
                <> display mean
  where
    runBench = timeout (minute 2) $
      benchmarkWith' (defaultConfig { verbosity = Quiet }) benchmarkable

    display :: Time Second -> Text
    display n = case n of
      (toUnit @Minute -> x) | x > minute 1 -> format x
      (toUnit @Second -> x) | x > sec 1 -> format x
      (toUnit @Millisecond -> x) | x > ms 1 -> format x
      (toUnit @Microsecond -> x) | x > mcs 1 -> format x
      (toUnit @Nanosecond -> x) | x > ns 1 -> format x
      _ -> format (toUnit @Picosecond n)

    format :: forall unit. KnownUnitName unit => Time unit -> Text
    format n =
      toText @String $ printf "%.4f%s"
        (realToFrac @RatioNat @Double $ unTime n)
        (unitNameVal @unit)

----------------------------------------------------------------------------
-- Generator
----------------------------------------------------------------------------

-- | Generates an @a@ and a @b@ and wraps them in a tuple.
genTuple2 :: MonadGen m => m a -> m b -> m (a, b)
genTuple2 = liftA2 (,)

-- | Construct a hedgehog generator from a generator relying on 'MonadRandom'.
--
-- This neither shrinks nor generates values in a reasonable order,
-- use only when such properties are justified for your type.
genRandom :: MonadGen m => (forall n. MonadRandom n => n a) -> m a
genRandom gen = evalRand gen . mkStdGen <$> Gen.prune Gen.enumBounded

-- | Run the given generator deterministically, by fixing its size and seed.
runGen :: HasCallStack => Range.Size -> Word64 -> Gen a -> a
runGen size seed genT =
  let tree = runGenT size (Seed.from seed) genT
      node = fromMaybe discardedErr $ runIdentity $ runMaybeT $ runTreeT tree
      discardedErr = error $
        "Generator could not produce a value for size "
        <> show (Range.unSize size) <> " and seed " <> show seed
  in  nodeValue node

----------------------------------------------------------------------------
-- Roundtrip
----------------------------------------------------------------------------

-- | This 'TestTree' contains a property based test for conversion from
-- some @x@ to some @y@ and back to @x@ (it should successfully return
-- the initial @x@).
roundtripTree
  :: forall x y err.
     ( Show x
     , Show y
     , Show err
     , Typeable x
     , Eq x
     , Eq err
     , HasCallStack
     )
  => Gen x
  -> (x -> y)
  -> (y -> Either err x)
  -> TestTree
roundtripTree genX xToY yToX = testProperty typeNameX prop
  where
    typeNameX = show $ typeRep (Proxy @x)
    prop :: Property
    prop = property $ do
      x <- forAll genX
      tripping x xToY yToX

-- | Test that one doc item goes before another doc item in generated
-- documentation.
assertGoesBefore
  :: forall d1 d2.
      (DocItem d1, DocItem d2)
  => Proxy d1 -> Proxy d2 -> Assertion
assertGoesBefore dp1 dp2 =
  unless (p1 < p2) $
    assertFailure $
      "Doc item " <> show (typeRep dp1) <> " with position " <> pretty p1 <> " \
      \goes before doc item " <> show (typeRep dp2) <> " with position " <> pretty p2
  where
    p1 = docItemPosition @d1
    p2 = docItemPosition @d2

-- | Test that one doc item goes before another doc item in generated
-- documentation.
goesBefore
  :: forall d1 d2.
      (DocItem d1, DocItem d2)
  => Proxy d1 -> Proxy d2 -> TestTree
goesBefore dp1 dp2 = testCase testName (assertGoesBefore dp1 dp2)
  where
  testName = "`" +| show @Text (typeRep dp1) |+ "` should come before `"
    +| show @Text (typeRep dp2) |+ "`"

----------------------------------------------------------------------------
-- Pretty-printing
----------------------------------------------------------------------------

formatValue :: forall t. SingI t => T.Value t -> Builder
formatValue v = "" +| build v |+ " of type " +| demote @t |+ ""

formatSomeValue :: (forall t. c t => SingI t) => SomeConstrainedValue c -> Builder
formatSomeValue = foldConstrained formatValue

-- | Derive a 'Show' instance for a type using a custom "show" function.
-- Note: the `show`n value is paren-wrapped iff it's a subexpression,
-- just as an ordinary `Show` would do.
data ShowWith a = ShowWith (a -> String) a

instance Eq a => Eq (ShowWith a) where
  ShowWith _ x == ShowWith _ y = x == y

instance Show (ShowWith a) where
  showsPrec d (ShowWith f a) = Text.Show.showParen (d > app_prec) $ Text.Show.showString $ f a
    where
      app_prec = 10

-- | Derive a 'Buildable' instance for a type using 'show'.
newtype Showing a = Showing a
  deriving stock Eq
  deriving newtype Show

instance Show a => Buildable (Showing a) where
  build (Showing a) = build (Debug.show @Text a)

----------------------------------------------------------------------------
-- Time
----------------------------------------------------------------------------

-- | Round the given time to the nearest whole number of the given unit,
-- not smaller than the given time.
--
-- @
-- ceilingUnit (sec 2.0) == sec 2
-- ceilingUnit (sec 2.1) == sec 3
-- ceilingUnit (sec 2.9) == sec 3
-- @
ceilingUnit :: forall (unit :: Rat) . Time unit -> Time unit
ceilingUnit = time . (% 1) . ceiling . unTime

-- | Converts the given time to a number with fixed-precision (in the given time unit).
--
-- @
-- timeToFixed (sec 1.234) == (1.2            :: Deci)
-- timeToFixed (sec 1.234) == (1.234          :: Milli)
-- timeToFixed (sec 1.234) == (1.234000000000 :: Pico)
-- @
timeToFixed
  :: forall precision unit
   . HasResolution precision
  => Time unit -> Fixed precision
timeToFixed = fromRational @(Fixed precision) . toRational . unTime

-- | Converts the given time to a 'NominalDiffTime'.
timeToNominalDiffTime :: KnownDivRat unit Second => Time unit -> NominalDiffTime
timeToNominalDiffTime =
  secondsToNominalDiffTime . timeToFixed . toUnit @Second

----------------------------------------------------------------------------
-- Bytes
----------------------------------------------------------------------------

stripOptional0x :: Text -> Text
stripOptional0x h = T.stripPrefix "0x" h ?: h

fromHex :: Text -> Either Text ByteString
fromHex hexText =
  let errMsg hex = "Invalid hex: \"" <> hex <> "\""
  in maybeToRight (errMsg hexText) . decodeHex . stripOptional0x $ hexText

parseAddressFromHex :: Text -> Either Text Address
parseAddressFromHex = fromHex >=> first pretty . parseAddressRaw

----------------------------------------------------------------------------
-- Traversals
----------------------------------------------------------------------------

{- | Version of 'mapM' generalized with 'L.each'.

Example:

> (addr1, addr2, addr3) <- mapEach newAddress ("test1", "test2", "test3")

This is more type-safe than simple 'mapM' since lists do not remember
their length in types.
-}
mapEach
  :: (L.Each s t a b, Applicative m)
  => (a -> m b) -> s -> m t
mapEach = L.traverseOf L.each

-- | Version of 'mapEach' with arguments flipped.
forEach
  :: (L.Each s t a b, Applicative m)
  => s -> (a -> m b) -> m t
forEach = L.forOf L.each