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

  -- * Compatibility
  TestName,
) where

import Control.Monad (forM)
import Control.Monad.Trans.State.Strict
import Control.Monad.IO.Class
import Data.List (isInfixOf)
import Data.Maybe (isJust, mapMaybe, fromJust)
import Data.String (fromString)
import Data.Time.Clock
import System.Environment
import System.Exit
import System.IO (hFlush, hPutStrLn, stdout, stderr)

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


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

type TestName = String

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

testGroupCollapse :: String -> [TestTree] -> TestTree
testGroupCollapse = Group True

-- | 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 (path++"/"++name)) 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` (path ++ "/" ++ name) = Just hp
      | otherwise = Nothing

computeMaxLen :: TestTree -> Int
computeMaxLen = go 0
  where
    go :: Int -> TestTree -> Int
    go indent (Group True name trees) = maximum (2*indent + length name : map (go (indent+1)) trees)
    go indent (Group False _ trees) = 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

newtype M a = M (StateT Stats IO a)
  deriving newtype (Functor, Applicative, Monad, MonadIO)

modifyStats :: (Stats -> Stats) -> M ()
modifyStats f = M (modify f)

data Options = Options
  { optsPattern :: String
  , optsHelp :: Bool }
  deriving (Show)

defaultOptions :: Options
defaultOptions = Options "" 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 (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'."]

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
      let M m = let ?maxlen = computeMaxLen tree in go 0 tree
      starttm <- getCurrentTime
      (success, stats) <- runStateT m initStats
      endtm <- getCurrentTime
      printStats stats (diffUTCTime endtm starttm)
      return (if isJust success then ExitSuccess else ExitFailure 1)
  where
    -- If all tests are successful, returns the number of output lines produced
    go :: (?maxlen :: Int) => Int -> TestTree -> M (Maybe Int)
    go indent (Group collapse name trees) = do
      liftIO $ putStrLn (replicate (2 * indent) ' ' ++ name)
      starttm <- liftIO getCurrentTime
      mlns <- fmap (fmap sum . sequence) . forM trees $ go (indent + 1)
      endtm <- liftIO getCurrentTime
      case mlns of
        Just lns | collapse -> 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 mlns
    go indent (Resource make cleanup fun) = do
      value <- liftIO make
      success <- go indent (fun value)
      liftIO $ cleanup value
      return success
    go indent (HP name (H.Property config test)) = do
      let thislen = 2*indent + length name
      liftIO $ putStr (replicate (2*indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ')
      liftIO $ hFlush stdout

      seed <- H.Seed.random

      starttm <- liftIO getCurrentTime
      report <- liftIO $ H.checkReport config 0 seed test (outputProgress (?maxlen + 2))
      endtm <- liftIO getCurrentTime

      liftIO $ printResult report (diffUTCTime endtm starttm)

      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)

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

printResult :: H.Report H.Result -> NominalDiffTime -> IO ()
printResult report timeTaken = do
  str <- H.renderResult H.EnableColor (Just (fromString "")) report
  if H.reportStatus report == H.OK
    then putStrLn ("\x1B[K" ++ str ++ prettyDuration False (realToFrac timeTaken))
    else putStrLn ("\x1B[K" ++ str)

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

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])