{-# LANGUAGE DeriveDataTypeable #-}
module Test.Hspec.Expectations.Pretty where

import Data.Dynamic (Typeable)
import Test.Hspec.Expectations (Expectation)
import Text.PrettyPrint.Free (pretty, renderPretty, displayS)
import System.Console.Terminfo.PrettyPrint (PrettyTerm(..))
import Control.Exception as E
import Control.DeepSeq (deepseq)
import Control.Monad (unless)

infix 1 `shouldBe`, `shouldSatisfy`, `shouldReturn`

data HspecFailure = HspecFailure String deriving (Typeable)
instance Show HspecFailure where show (HspecFailure msg) = msg
instance Exception HspecFailure

assertFailure :: String -- ^ A message that is displayed with the assertion failure 
              -> Expectation
assertFailure msg = msg `deepseq` E.throwIO (HspecFailure msg)

-- | Asserts that the specified actual value is equal to the expected value.
-- The output message will contain the prefix, the expected value, and the 
-- actual value.
--  
-- If the prefix is the empty string (i.e., @\"\"@), then the prefix is omitted
-- and only the expected and actual values are output.
assertEqual :: (Eq a, PrettyTerm a) => String -- ^ The message prefix 
                              -> a      -- ^ The expected value 
                              -> a      -- ^ The actual value
                              -> Expectation
assertEqual preface expected actual =
  unless (actual == expected) $ assertFailure msg
 where
   msg = (if null preface then "" else preface ++ "\n") ++
           "expected: " ++ render (pretty expected) ++
        "\n  but got: " ++ render (pretty actual)
   render x = displayS (renderPretty 0.4 80 x) ""
             

-- |
-- @actual \`shouldBe\` expected@ sets the expectation that @actual@ is equal
-- to @expected@ (this is just an alias for `@?=`).
shouldBe :: (Show a, PrettyTerm a, Eq a) => a -> a -> Expectation
actual `shouldBe` expected = assertEqual "" expected actual

-- |
-- @v \`shouldSatisfy\` p@ sets the expectation that @p v@ is @True@.
shouldSatisfy :: (Show a, PrettyTerm a) => a -> (a -> Bool) -> Expectation
v `shouldSatisfy` p = assertBool ("predicate failed on: " ++ show v) (p v)
  where assertBool msg b = unless b (assertFailure msg)

-- |
-- @action \`shouldReturn\` expected@ sets the expectation that @action@
-- returns @expected@.
shouldReturn :: (Show a, PrettyTerm a, Eq a) => IO a -> a -> Expectation
action `shouldReturn` expected = action >>= (`shouldBe` expected)