    {-# LANGUAGE DeriveDataTypeable #-}
    module Test.Tasty.Ingredients.Rerun (rerunningTests) where

    import Prelude hiding (filter)

    import Control.Applicative
    import Control.Arrow ((>>>))
    import Control.Exception (throwIO)
    import Control.Monad (guard, when)
    import Control.Monad.Trans.Class (lift)
    import Data.Char (isSpace)
    import Data.Foldable (asum)
    import Data.List.Split (endBy)
    import Data.Maybe (fromMaybe)
    import Data.Monoid (mconcat)
    import Data.Proxy (Proxy(..))
    import Data.Semigroup.Applicative (Traversal(..))
    import Data.Tagged (Tagged(..), untag)
    import Data.Typeable (Typeable)
    import System.IO.Error (catchIOError, isDoesNotExistError)

    import qualified Control.Concurrent.STM as STM
    import qualified Control.Monad.State as State
    import qualified Data.Functor.Compose as Functor
    import qualified Data.IntMap as IntMap
    import qualified Data.Map.Strict as Map
    import qualified Data.Set as Set
    import qualified Options.Applicative as OptParse
    import qualified Test.Tasty as Tasty
    import qualified Test.Tasty.Options as Tasty
    import qualified Test.Tasty.Providers as Tasty
    import qualified Test.Tasty.Runners as Tasty

    --------------------------------------------------------------------------------
    newtype RerunLogFile = RerunLogFile FilePath
      deriving (Typeable)

    instance Tasty.IsOption RerunLogFile where
      optionName = Tagged "rerun-log-file"
      optionHelp = Tagged "The path to which rerun's state file should be saved"
      defaultValue = RerunLogFile ".tasty-rerun-log"
      parseValue = Just . RerunLogFile

    --------------------------------------------------------------------------------
    newtype UpdateLog = UpdateLog Bool
      deriving (Typeable)

    instance Tasty.IsOption UpdateLog where
      optionName = Tagged "rerun-update"

      optionHelp = Tagged "If present the log file will be updated, otherwise it \
                          \will be left unchanged"

      defaultValue = UpdateLog False

      parseValue = Just . UpdateLog . const True

      optionCLParser = fmap UpdateLog $ OptParse.switch $ mconcat
        [ OptParse.long name
        , OptParse.help helpString
        ]
        where
        name = untag (Tasty.optionName :: Tagged UpdateLog String)
        helpString = untag (Tasty.optionHelp :: Tagged UpdateLog String)

    --------------------------------------------------------------------------------
    data Filter = Failures | Exceptions | New | Successful
      deriving (Eq, Ord)

    parseFilter :: String -> Maybe Filter
    parseFilter "failures" = Just Failures
    parseFilter "exceptions" = Just Exceptions
    parseFilter "new" = Just New
    parseFilter "successful" = Just Successful
    parseFilter _ = Nothing

    --------------------------------------------------------------------------------
    everything :: [Filter]
    everything = [Failures, Exceptions, New, Successful]

    --------------------------------------------------------------------------------
    newtype FilterOption = FilterOption (Set.Set Filter)
      deriving (Typeable)

    instance Tasty.IsOption FilterOption where
      optionName = Tagged "rerun-filter"

      optionHelp = Tagged "A comma separated list to specify which tests to run when\
                          \ comparing against previous test runs. Valid options \
                          \are: everything, failures, exceptions, new"

      defaultValue = FilterOption (Set.fromList everything)

      parseValue =
        fmap (FilterOption . Set.fromList) . mapM (parseFilter . trim) . endBy ","
        where trim = reverse . dropWhile isSpace . reverse . dropWhile isSpace

    --------------------------------------------------------------------------------
    data TestResult = Completed Bool | ThrewException
      deriving (Read, Show)

    --------------------------------------------------------------------------------
    -- | This 'Tasty.Ingredient' transformer adds various @--rerun@ options to your
    -- test program. These flags add stateful execution of your test suite, allowing
    -- you to rerun only tests that are failing from the previous run, or tests that
    -- that have been added since the last test ran, once the 'Tasty.TestTree' has
    -- been filtered.
    --
    -- The input list of 'Tasty.Ingredient's specifies the 'Tasty.Ingredients's that
    -- will actually work with the filtered 'Tasty.TestTree'. Normally, you'll want
    -- at least 'Tasty.Test.Runners.consoleTestReporter'.
    rerunningTests :: [Tasty.Ingredient] -> Tasty.Ingredient
    rerunningTests ingredients =
      Tasty.TestManager (rerunOptions ++ existingOptions) $
        \options testTree -> Just $ do
          let RerunLogFile stateFile = Tasty.lookupOption options
          let UpdateLog updateLog = Tasty.lookupOption options
          let FilterOption filter = Tasty.lookupOption options

          testTree' <- maybe (Just testTree) (filterTestTree options testTree filter)
                         <$> tryLoadStateFrom stateFile

          case testTree' of
            -- We filtered the test tree down to 0 tests
            Nothing -> return True

            -- There are tests to run so we try and find an Ingredient to run the
            -- tests
            Just filteredTestTree -> do
              let tryAndRun (Tasty.TestReporter _ f) = do
                    runner <- f options filteredTestTree
                    return $ do
                      statusMap <- Tasty.launchTestTree options filteredTestTree
                      let getTestResults =
                            fmap getConst $
                            flip State.evalStateT 0 $
                            Functor.getCompose $
                            getTraversal $
                            Tasty.foldTestTree (observeResults statusMap)
                                               options filteredTestTree
                      outcome <- runner statusMap
                      when updateLog (saveStateTo stateFile getTestResults)
                      return outcome

                  tryAndRun (Tasty.TestManager _ f) =
                    f options filteredTestTree

              case asum (map tryAndRun ingredients) of
                -- No Ingredients chose to run the tests, we should really return
                -- Nothing, but we've already commited to run by the act of
                -- filtering the TestTree.
                Nothing -> return False

                -- Otherwise, an Ingredient did choose to run the tests, so we
                -- simply run the above constructed IO action.
                Just e -> e

      where
      existingOptions = flip concatMap ingredients $ \ingredient ->
        case ingredient of
          Tasty.TestReporter options _ -> options
          Tasty.TestManager options _ -> options

      rerunOptions = [ Tasty.Option (Proxy :: Proxy RerunLogFile)
                     , Tasty.Option (Proxy :: Proxy UpdateLog)
                     , Tasty.Option (Proxy :: Proxy FilterOption)
                     ]

      ------------------------------------------------------------------------------
      filterTestTree options testTree filter lastRecord =
        let foldSingle _ name t = \prefix ->
              let requiredFilter = case Map.lookup (prefix ++ [name]) lastRecord of
                    Just (Completed False) -> Failures
                    Just ThrewException -> Exceptions
                    Just (Completed True) -> Successful
                    Nothing -> New

              in do guard (requiredFilter `Set.member` filter)
                    return (Tasty.SingleTest name t)

            foldGroup name tests = \prefix ->
              [ Tasty.testGroup name (tests (prefix ++ [name])) ]

            foldResource rSpec k = \prefix ->
              let peek = k (error "Resources unavailable during test tree filtering")
              in case peek prefix of
                   [x] -> [ Tasty.WithResource rSpec (\io -> head $ k io prefix) ]
                   []  -> []
                   _   -> error "Resource claims to initialize multiple tests. \
                                \Please report this as a bug."

            treeFold = Tasty.TreeFold { Tasty.foldSingle = foldSingle
                                      , Tasty.foldGroup = foldGroup
                                      , Tasty.foldResource = foldResource
                                      }

        in case Tasty.foldTestTree treeFold options testTree [] of
             [t] -> Just t
             [] -> Nothing

             -- This state is impossible as a TestTree is either a single test
             -- or a test group. Test groups are folded into a single list element
             -- and single tests to 0-or-1. Thus I believe this state is impossible.
             _ ->
                error "tasty-rerun found multiple tests when one was expected. \
                      \If you can produce this error, please report this as a bug!"

      ------------------------------------------------------------------------------
      tryLoadStateFrom filePath = do
        fileContents <- (Just <$> readFile filePath)
                          `catchIOError` (\e -> if isDoesNotExistError e
                                                  then return Nothing
                                                  else ioError e)
        return (read <$> fileContents)

      ------------------------------------------------------------------------------
      saveStateTo filePath getTestResults =
        getTestResults >>= (show >>> writeFile filePath)

      ------------------------------------------------------------------------------
      observeResults statusMap =
        let foldSingle _ name _ = Traversal $ Functor.Compose $ do
              i <- State.get

              status <- lift $ STM.atomically $ do
                status <- lookupStatus i
                case status of
                  Tasty.Done result ->
                    return (Completed (Tasty.resultSuccessful result))

                  Tasty.Exception _ -> return ThrewException

                  _ -> STM.retry

              Const (Map.singleton [name] status) <$ State.modify (+ 1)

            foldGroup name children = Traversal $ Functor.Compose $ do
              Const soFar <- Functor.getCompose $ getTraversal children
              pure $ Const (Map.mapKeys (name :) soFar)

        in Tasty.trivialFold
          { Tasty.foldSingle = foldSingle
          , Tasty.foldGroup = foldGroup
          }

        where
        lookupStatus i = STM.readTVar $
          fromMaybe (error "Attempted to lookup test by index outside bounds")
                    (IntMap.lookup i statusMap)