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

  -- * Compatibility
  TestName,
) where

import Control.Monad (forM_)
import Control.Monad.Trans.State.Strict
import Control.Monad.IO.Class
import Data.String (fromString)
import Data.Time.Clock
import System.Exit
import System.IO (hFlush, stdout)

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 String [TestTree]
  | forall a. Resource (IO a) (a -> IO ()) (a -> TestTree)
  | HP String H.Property

type TestName = String

testGroup :: String -> [TestTree] -> TestTree
testGroup = 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

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

runTests :: TestTree -> IO ExitCode
runTests = \tree -> do
  let M m = let ?maxlen = computeMaxLen tree in go 0 tree
  starttm <- getCurrentTime
  stats <- execStateT m initStats
  endtm <- getCurrentTime
  printStats stats (diffUTCTime endtm starttm)
  return (if statsOK stats == statsTotal stats
            then ExitSuccess else ExitFailure 1)
  where
    go :: (?maxlen :: Int) => Int -> TestTree -> M ()
    go indent (Group name trees) = do
      liftIO $ putStrLn (replicate (2 * indent) ' ' ++ name)
      forM_ trees $ go (indent + 1)
    go indent (Resource make cleanup fun) = do
      value <- liftIO make
      go indent (fun value)
      liftIO $ cleanup value
      return ()
    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) name)
      endtm <- liftIO getCurrentTime

      liftIO $ printResult name report (diffUTCTime endtm starttm)

      modifyStats $ \stats -> stats { statsOK = fromEnum (H.reportStatus report == H.OK) + statsOK stats
                                    , statsTotal = 1 + statsTotal stats }

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

printResult :: String -> H.Report H.Result -> NominalDiffTime -> IO ()
printResult name report timeTaken = do
  str <- H.renderResult H.EnableColor (Just (fromString name)) report
  if timeTaken >= 0.01
    then putStrLn ("\x1B[K" ++ str ++ " (" ++ prettySeconds (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 (" ++ prettySeconds (realToFrac timeTaken) ++ ")\x1B[0m"
  | otherwise =
      let nfailed = statsTotal stats - statsOK stats
      in putStrLn $ "\x1B[31mFailed " ++ show nfailed ++ " out of " ++ show (statsTotal stats) ++
                    " tests (" ++ prettySeconds (realToFrac timeTaken) ++ ")\x1B[0m"

prettySeconds :: Double -> String
prettySeconds 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"

defaultMain :: TestTree -> IO ()
defaultMain tree = runTests tree >>= exitWith