{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE TupleSections #-}
module Test.Framework (
  TestTree,
  testGroup,
  groupSetCollapse, groupSetSequential,
  testProperty,
  withResource,
  withResource',
  runTests,
  defaultMain,
  Options(..),

  -- * Compatibility
  TestName,
) where

import Control.Concurrent
import Control.Exception (finally)
import Control.Monad (forM, when, replicateM)
import Control.Monad.IO.Class
import Data.IORef
import Data.List (isInfixOf, intercalate)
import Data.Maybe (isJust, mapMaybe, fromJust)
import Data.String (fromString)
import Data.Time.Clock
import GHC.Conc (getNumProcessors)
import System.Console.Concurrent
import System.Console.Regions
import System.Environment
import System.Exit
import System.IO (hFlush, hPutStrLn, stdout, stderr, hIsTerminalDevice)
import Text.Read (readMaybe)

import qualified Hedgehog as H
import qualified Hedgehog.Internal.Config as H
import qualified Hedgehog.Internal.Property as H
import qualified Hedgehog.Internal.Report as H
import qualified Hedgehog.Internal.Runner as H
import qualified Hedgehog.Internal.Seed as H.Seed


type TestName = String

data TestTree
  = Group GroupOpts String [TestTree]
  | forall a. Resource (IO a) (a -> IO ()) (a -> TestTree)
  | HP String H.Property

data GroupOpts = GroupOpts
  { goCollapse :: Bool
  , goSequential :: Bool }
  deriving (Show)

defaultGroupOpts :: GroupOpts
defaultGroupOpts = GroupOpts False False

testGroup :: String -> [TestTree] -> TestTree
testGroup = Group defaultGroupOpts

groupSetCollapse :: TestTree -> TestTree
groupSetCollapse (Group opts name trees) = Group opts { goCollapse = True } name trees
groupSetCollapse _ = error "groupSetCollapse: not called on a Group"

groupSetSequential :: TestTree -> TestTree
groupSetSequential (Group opts name trees) = Group opts { goSequential = True } name trees
groupSetSequential _ = error "groupSetSequential: not called on a Group"

-- | The @a -> TestTree@ function must use the @a@ only inside properties: when
-- not actually running properties, it will be passed 'undefined'.
withResource :: IO a -> (a -> IO ()) -> (a -> TestTree) -> TestTree
withResource = Resource

-- | Same caveats as 'withResource'.
withResource' :: IO a -> (a -> TestTree) -> TestTree
withResource' make fun = withResource make (\_ -> return ()) fun

testProperty :: String -> H.Property -> TestTree
testProperty = HP

filterTree :: Options -> TestTree -> Maybe TestTree
filterTree (Options { optsPattern = pat }) = go []
  where
    go path (Group collapse name trees) =
      case mapMaybe (go (name:path)) trees of
        [] -> Nothing
        trees' -> Just (Group collapse name trees')
    go path (Resource make free fun) =
      case go path (fun undefined) of
        Nothing -> Nothing
        Just _ -> Just $ Resource make free (fromJust . go path . fun)
    go path hp@(HP name _)
      | pat `isInfixOf` renderPath (name:path) = Just hp
      | otherwise = Nothing

    renderPath comps = "^" ++ intercalate "/" (reverse comps) ++ "$"

computeMaxLen :: TestTree -> Int
computeMaxLen = go 0
  where
    go :: Int -> TestTree -> Int
    go indent (Group opts name trees)
      -- If we collapse, the name of the group gets prefixed before the final status message after collapsing.
      | goCollapse opts = maximum (2*indent + length name : map (go (indent+1)) trees)
      -- If we don't collapse, the group name does get printed but without any status message, so it doesn't need to get accounted for in maxlen.
      | otherwise = maximum (0 : map (go (indent+1)) trees)
    go indent (Resource _ _ fun) = go indent (fun undefined)
    go indent (HP name _) = 2 * indent + length name

data Stats = Stats
  { statsOK :: Int
  , statsTotal :: Int }
  deriving (Show)

initStats :: Stats
initStats = Stats 0 0

modifyStats :: (?stats :: IORef Stats) => (Stats -> Stats) -> IO ()
modifyStats f = atomicModifyIORef' ?stats (\s -> (f s, ()))

data Options = Options
  { optsPattern :: String
  , optsHelp :: Bool
  , optsHedgehogReplay :: Maybe (H.Skip, H.Seed)
  , optsHedgehogShrinks :: Maybe Int
  , optsParallel :: Bool
  }
  deriving (Show)

defaultOptions :: Options
defaultOptions = Options "" False Nothing Nothing False

parseOptions :: [String] -> Options -> Either String Options
parseOptions [] opts = pure opts
parseOptions ("-h":args) opts = parseOptions args opts { optsHelp = True }
parseOptions ("--help":args) opts = parseOptions args opts { optsHelp = True }
parseOptions ("-p":arg:args) opts
  | optsPattern opts == "" = parseOptions args opts { optsPattern = arg }
  | otherwise = Left "Multiple '-p' arguments given"
parseOptions ("--hedgehog-replay":arg:args) opts
  | Nothing <- optsHedgehogReplay opts =
      let parsed = do
            (skipstr, ' ' : seedstr) <- return $ span (/= ' ') arg
            (,) <$> H.skipDecompress skipstr <*> readMaybe seedstr
      in case parsed of
           Just res -> parseOptions args opts { optsHedgehogReplay = Just res }
           Nothing -> Left "Invalid argument to '--hedgehog-replay'"
  | otherwise = Left "Multiple '--hedgehog-replay' arguments given"
parseOptions ("--hedgehog-shrinks":arg:args) opts =
  case readMaybe arg of
    Just n -> parseOptions args opts { optsHedgehogShrinks = Just n }
    Nothing -> Left "Invalid argument to '--hedgehog-shrinks'"
parseOptions ("--parallel":args) opts = parseOptions args opts { optsParallel = True }
parseOptions (arg:_) _ = Left $ "Unrecognised argument: '" ++ arg ++ "'"

printUsage :: IO ()
printUsage = do
  progname <- getProgName
  putStr $ unlines
    ["Usage: " ++ progname ++ " [options]"
    ,"Options:"
    ,"  -h / --help    Show this help"
    ,"  -p PATTERN     Only tests whose path contains PATTERN are run. The path of a"
    ,"                 test looks like: '^group1/group2/testname$'."
    ,"  --hedgehog-replay '{skip} {seed}'"
    ,"                 Skip to a particular generated Hedgehog test. Should be used"
    ,"                 with -p. Overrides 'propertySkip' in 'PropertyConfig' if set."
    ,"  --hedgehog-shrinks NUM"
    ,"                 Limit the number of shrinking steps."
    ,"  --parallel     Run tests in parallel."]

defaultMain :: TestTree -> IO ()
defaultMain tree = do
  args <- getArgs
  case parseOptions args defaultOptions of
    Left err -> die err
    Right opts
      | optsHelp opts -> printUsage >> exitSuccess
      | otherwise -> runTests opts tree >>= exitWith

runTests :: Options -> TestTree -> IO ExitCode
runTests options = \tree' ->
  case filterTree options tree' of
    Nothing -> do hPutStrLn stderr "No tests matched the given pattern."
                  return (ExitFailure 1)
    Just tree -> do
      isterm <- hIsTerminalDevice stdout
      starttm <- getCurrentTime
      statsRef <- newIORef initStats
      success <- let ?stats = statsRef
                     ?options = options
                     ?maxlen = computeMaxLen tree
                     ?istty = isterm
                 in if optsParallel options
                      then do nproc <- getNumProcessors
                              setNumCapabilities nproc
                              displayConsoleRegions $
                                withWorkerPool nproc $ \ pool ->
                                  let ?pool = pool in runTreePar Nothing 0 id tree
                      else isJust <$> runTreeSeq 0 id tree
      stats <- readIORef statsRef
      endtm <- getCurrentTime
      let ?istty = isterm in printStats stats (diffUTCTime endtm starttm)
      return (if success then ExitSuccess else ExitFailure 1)

-- If all tests are successful, returns the number of output lines produced
runTreeSeq :: (?options :: Options,  ?stats :: IORef Stats,?maxlen :: Int, ?istty :: Bool)
           => Int -> (String -> String) -> TestTree -> IO (Maybe Int)
runTreeSeq indent path (Group groupOpts name trees) = do
  liftIO $ putStrLn (replicate (2 * indent) ' ' ++ name) >> hFlush stdout
  starttm <- liftIO getCurrentTime
  mlns <- fmap (fmap sum . sequence) . forM trees $
            runTreeSeq (indent + 1) (path . (name++) . ('/':))
  endtm <- liftIO getCurrentTime
  case mlns of
    Just lns | goCollapse groupOpts, ?istty -> do
      let thislen = 2*indent + length name
      liftIO $ putStrLn $ concat (replicate (lns+1) "\x1B[A\x1B[2K") ++ "\x1B[G" ++
                          replicate (2 * indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ' ++
                          "\x1B[32mOK\x1B[0m" ++
                          prettyDuration False (realToFrac (diffUTCTime endtm starttm))
      return (Just 1)
    _ -> return ((+1) <$> mlns)
runTreeSeq indent path (Resource make cleanup fun) = do
  value <- liftIO make
  success <- runTreeSeq indent path (fun value)
  liftIO $ cleanup value
  return success
runTreeSeq indent path (HP name prop) =
  runHP (\prefix -> when ?istty $ putStr prefix >> hFlush stdout)
        (\_ -> outputProgress (?maxlen + 2))
        (\prefix rendered -> putStrLn ((if ?istty then "\x1B[K" else prefix) ++ rendered) >> hFlush stdout)
        indent path name prop

-- Assumes it's run within displayConsoleRegions.
runTreePar :: (?options :: Options, ?stats :: IORef Stats, ?pool :: WorkerPool, ?maxlen :: Int, ?istty :: Bool)
           => Maybe (ConsoleRegion, String) -> Int -> (String -> String) -> TestTree -> IO Bool
runTreePar mregctx indent path (Group groupOpts name trees) = do
  let run reg regPrefix sequential = do
        setConsoleRegion reg name
        starttm <- liftIO getCurrentTime
        success <- fmap and . poolRunList ?pool . flip map trees $
          runTreeParSub reg (name ++ " > ") (indent + 1) (path . (name++) . ('/':))
        endtm <- liftIO getCurrentTime

        let thislen = 2*indent + length name
        finishConsoleRegion reg $
          replicate (2 * indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ' ++
          ansi "\x1B[32mOK\x1B[0m" ++
          prettyDuration False (realToFrac (diffUTCTime endtm starttm))
        return success

  case (mregctx, goSequential groupOpts) of
    (Nothing, True) -> do
      outputConcurrent (replicate (2 * indent) ' ' ++ name ++ "\n")
      fmap and . forM trees $
        runTreePar Nothing (indent + 1) (path . (name++) . ('/':))
    (_, False) -> do
      regPrefix <- case mregctx of
        Just (reg, regPrefix) -> do
          setConsoleRegion reg (regPrefix ++ name)
          return regPrefix
        Nothing -> return ""

      starttm <- liftIO getCurrentTime
      success <- fmap and . poolRunList ?pool . flip map trees $ \tree ->
        withConsoleRegion Linear $ \reg ->
          runTreePar (Just (reg, regPrefix ++ name ++ " > "))
                     (indent + 1) (path . (name++) . ('/':)) tree
      endtm <- liftIO getCurrentTime

      let thislen = 2*indent + length name
      finishConsoleRegion reg $
        replicate (2 * indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ' ++
        ansi "\x1B[32mOK\x1B[0m" ++
        prettyDuration False (realToFrac (diffUTCTime endtm starttm))
      return success
    (Just (reg, regPrefix), sequential) ->
      run reg regPrefix sequential

runTreePar mregctx indent path (Resource make cleanup fun) = do
  value <- liftIO make
  success <- runTreePar mregctx indent path (fun value)
  liftIO $ cleanup value
  return success

runTreePar mregctx indent path (HP name prop) =
  let run reg regPrefix =
        isJust <$>
          runHP (\prefix -> setConsoleRegion reg (regPrefix ++ prefix))
                (\prefix -> outputProgressPar reg (regPrefix ++ prefix))
                (\prefix rendered -> finishConsoleRegion reg (regPrefix ++ prefix ++ rendered ++ "\n"))
                indent path name prop
  in case mregctx of
       Nothing -> withConsoleRegion Linear $ \reg -> run reg ""
       Just (reg, regPrefix) -> run reg regPrefix

-- Sequential subcomputation in a parallel environment
runTreeParSub :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int, ?istty :: Bool)
              => ConsoleRegion -> String -> Int -> (String -> String) -> TestTree -> IO Bool
runTreeParSub region regPrefix indent path (Group _ name trees) =
  fmap and . forM trees $
    runTreeParSub region (regPrefix ++ name ++ " > ") (indent + 1) (path . (name++) . ('/':))

runTreeParSub region regPrefix indent path (Resource make cleanup fun) = do
  value <- liftIO make
  success <- runTreeParSub region regPrefix indent path (fun value)
  liftIO $ cleanup value
  return success

runTreeParSub region regPrefix indent path (HP name prop) = do
  isJust <$>
    runHP (\prefix -> setConsoleRegion region (regPrefix ++ prefix))
          (\prefix -> outputProgressPar region (regPrefix ++ prefix))
          (\prefix rendered -> finishConsoleRegion region (regPrefix ++ prefix ++ rendered))
          indent path name prop

runHP :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int, ?istty :: Bool)
      => (String -> IO ())
      -> (String -> H.Report H.Progress -> IO ())
      -> (String -> String -> IO ())
      -> Int -> (String -> String)
      -> String -> H.Property -> IO (Maybe Int)
runHP prefixPrinter progressPrinter resultPrinter indent path name (H.Property config test) = do
  let thislen = 2*indent + length name
  let outputPrefix = replicate (2*indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' '
  liftIO $ prefixPrinter outputPrefix

  let (config', seedfun) = applyHedgehogOptions ?options config
  seed <- seedfun

  starttm <- liftIO getCurrentTime
  report <- liftIO $ H.checkReport config' 0 seed test (progressPrinter outputPrefix)
  endtm <- liftIO getCurrentTime

  rendered <- liftIO $ renderResult report (path name) (diffUTCTime endtm starttm)
  liftIO $ resultPrinter outputPrefix rendered

  let ok = H.reportStatus report == H.OK
  modifyStats $ \stats -> stats { statsOK = fromEnum ok + statsOK stats
                                , statsTotal = 1 + statsTotal stats }
  return (if ok then Just 1 else Nothing)

applyHedgehogOptions :: MonadIO m => Options -> H.PropertyConfig -> (H.PropertyConfig, m H.Seed)
applyHedgehogOptions opts config0 =
  let (f1, seedfun) = case optsHedgehogReplay opts of
                        Just (skip, seed) -> (\c -> c { H.propertySkip = Just skip }, return seed)
                        Nothing -> (id, H.Seed.random)
      f2 = case optsHedgehogShrinks opts of
             Just n -> \c -> c { H.propertyShrinkLimit = H.ShrinkLimit n }
             Nothing -> id
  in (f2 . f1 $ config0, seedfun)

outputProgress :: (?istty :: Bool) => Int -> H.Report H.Progress -> IO ()
outputProgress indent report
  | ?istty = do
      str <- H.renderProgress H.EnableColor (Just (fromString "")) report
      putStr (replace '\n' "  " str ++ "\x1B[" ++ show (indent+1) ++ "G")
      hFlush stdout
  | otherwise = return ()

outputProgressPar :: ConsoleRegion -> String -> H.Report H.Progress -> IO ()
outputProgressPar region prefix report = do
  str <- H.renderProgress H.EnableColor (Just (fromString "")) report
  setConsoleRegion region (prefix ++ replace '\n' "  " str)

renderResult :: (?istty :: Bool) => H.Report H.Result -> String -> NominalDiffTime -> IO String
renderResult report path timeTaken = do
  str <- H.renderResult H.EnableColor (Just (fromString "")) report
  case H.reportStatus report of
    H.OK -> return (str ++ prettyDuration False (realToFrac timeTaken))
    H.Failed failure -> do
      let H.Report { H.reportTests = count, H.reportDiscards = discards } = report
          replayInfo = H.skipCompress (H.SkipToShrink count discards (H.failureShrinkPath failure)) ++
                       " " ++ show (H.reportSeed report)
          suffix = "\n    Flags to reproduce: `-p '" ++ path ++ "' --hedgehog-replay '" ++ replayInfo ++ "'`"
      return (str ++ suffix)
    _ -> return str

printStats :: (?istty :: Bool) => Stats -> NominalDiffTime -> IO ()
printStats stats timeTaken
  | statsOK stats == statsTotal stats = do
      putStrLn $ ansi "\x1B[32m" ++ "All " ++ show (statsTotal stats) ++
                 " tests passed." ++ prettyDuration True (realToFrac timeTaken) ++ ansi "\x1B[0m"
  | otherwise =
      let nfailed = statsTotal stats - statsOK stats
      in putStrLn $ ansi "\x1B[31m" ++ "Failed " ++ show nfailed ++ " out of " ++ show (statsTotal stats) ++
                    " tests." ++ prettyDuration True (realToFrac timeTaken) ++ ansi "\x1B[0m"

data WorkerPool = WorkerPool (Chan (Maybe PoolJob)) [ThreadId]
data PoolJob = forall a. PoolJob (IO a) (MVar a)

withWorkerPool :: Int -> (WorkerPool -> IO a) -> IO a
withWorkerPool numWorkers k = do
  chan <- newChan
  pool <- WorkerPool chan <$> forM [0..numWorkers-1] (\i -> forkOn i (worker i chan))
  k pool `finally` replicateM numWorkers (writeChan chan Nothing)
  where
    worker :: Int -> Chan (Maybe PoolJob) -> IO ()
    worker idx chan = do
      mjob <- readChan chan
      case mjob of
        Just (PoolJob action mvar) -> do
          outputConcurrent $ "[" ++ show idx ++ "] got job\n"
          action >>= putMVar mvar
          worker idx chan
        Nothing -> return ()

poolSubmit :: WorkerPool -> IO a -> MVar a -> IO ()
poolSubmit (WorkerPool chan _) action mvar = writeChan chan (Just (PoolJob action mvar))

poolRunList :: WorkerPool -> [IO a] -> IO [a]
poolRunList pool actions = do
  vars <- forM actions $ \act -> do
    var <- newEmptyMVar
    poolSubmit pool act var
    return var
  mapM takeMVar vars


prettyDuration :: Bool -> Double -> String
prettyDuration False x | x < 0.5 = ""
prettyDuration _ x =
  let str = show (round (x * 100) :: Int)
      str' = replicate (3 - length str) '0' ++ str
      (pre, post) = splitAt (length str' - 2) str'
  in " (" ++ pre ++ "." ++ post ++ "s)"

replace :: Eq a => a -> [a] -> [a] -> [a]
replace x ys = concatMap (\y -> if y == x then ys else [y])

ansi :: (?istty :: Bool) => String -> String
ansi | ?istty = id
     | otherwise = const ""