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

  -- * Compatibility
  TestName,
) where

import Control.Concurrent (setNumCapabilities, forkIO, killThread, forkOn)
import Control.Concurrent.MVar
import Control.Concurrent.STM
import Control.Exception (finally)
import Control.Monad (forM, when, forM_, replicateM_)
import Control.Monad.IO.Class
import Data.Colour.SRGB.Linear
import Data.IORef
import Data.List (isInfixOf, intercalate)
import Data.Maybe (isJust, mapMaybe, fromJust)
import Data.PQueue.Prio.Min qualified as PQ
import Data.String (fromString)
import Data.Time.Clock
import GHC.Conc (getNumProcessors)
import System.Console.ANSI qualified as ANSI
import System.Console.Concurrent (outputConcurrent)
import System.Console.Regions
import System.Environment (getArgs, getProgName)
import System.Exit
import System.IO (hFlush, hPutStrLn, stdout, stderr, hIsTerminalDevice)
import Text.Read (readMaybe)

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


type TestName = String

data TestTree
  = Group GroupOpts String [TestTree]
  | forall a. Resource String (IO a) (a -> IO ()) (a -> TestTree)
      -- ^ Name is not specified by user, but inherited from the tree below
  | HP String H.Property

-- Not exported because a Resource is not supposed to have a name in the first place
treeName :: TestTree -> String
treeName (Group _ name _) = name
treeName (Resource name _ _ _) = name
treeName (HP name _) = name

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

defaultGroupOpts :: GroupOpts
defaultGroupOpts = GroupOpts 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"

-- | The @a -> TestTree@ function must use the @a@ only inside properties: the
-- functoin will be passed 'undefined' when exploring the test tree (without
-- running properties).
withResource :: IO a -> (a -> IO ()) -> (a -> TestTree) -> TestTree
withResource make cleanup fun = Resource (treeName (fun undefined)) make cleanup fun

-- | 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 opts name trees) =
      case mapMaybe (go (name:path)) trees of
        [] -> Nothing
        trees' -> Just (Group opts name trees')
    go path (Resource inhname make free fun) =
      case go path (fun undefined) of
        Nothing -> Nothing
        Just _ -> Just $ Resource inhname 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
  , optsVerbose :: Bool
  }
  deriving (Show)

defaultOptions :: Options
defaultOptions = Options "" False Nothing Nothing False 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 ("--verbose":args) opts = parseOptions args opts { optsVerbose = 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."
    ,"  --verbose      Currently only has an effect with --parallel. Also shows OK"
    ,"                 and timing for test groups, not only individual tests."]

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 -> do
                         let ?pool = pool
                         successVar <- newEmptyMVar
                         runTreePar Nothing [] [] tree successVar
                         readMVar successVar
             else isJust <$> runTreeSeq 0 [] tree
      stats <- readIORef statsRef
      endtm <- getCurrentTime
      let ?istty = isterm in printStats stats (diffUTCTime endtm starttm)
      return (if success then ExitSuccess else ExitFailure 1)

-- | Returns when all jobs in this tree have been scheduled. When all jobs are
-- done, the outvar is filled with whether all tests in this tree were
-- successful.
-- The mparregion is the parent region to take over, if any. Having a parent
-- region to take over implies that we are currently executing in a worker and
-- can hence run blocking user code directly in the current thread.
runTreePar :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int, ?istty :: Bool, ?pool :: WorkerPool [Int])
           => Maybe ConsoleRegion -> [Int] -> [String] -> TestTree -> MVar Bool -> IO ()
-- TODO: handle collapse somehow?
runTreePar mparregion revidxlist revpath (Group _ name trees) outvar = do
  let path = intercalate "/" (reverse (name : revpath))
  -- outputConcurrent $ "! " ++ path ++ ": Started\n"

  -- If we have exactly one child and we're currently running in a worker, we can continue doing so
  mparregion2 <- case trees of
    [_] -> return mparregion
    _ -> do  -- If not, we have to close the region (and implicitly relinquish the worker job)
      forM_ mparregion closeConsoleRegion
      return Nothing

  starttm <- getCurrentTime
  suboutvars <- forM (zip trees [0..]) $ \(tree, idx) -> do
    suboutvar <- newEmptyMVar
    runTreePar mparregion2 (idx : revidxlist) (name : revpath) tree suboutvar
    return suboutvar

  -- outputConcurrent $ "! " ++ path ++ ": Scheduled all\n"

  -- If we took over the parent region then this readMVar will resolve
  -- immediately and the forkIO would be unnecessary. Meh.
  _ <- forkIO $ do
    success <- and <$> traverse readMVar suboutvars
    endtm <- getCurrentTime
    -- outputConcurrent $ "! " ++ path ++ ": Done\n"
    if success && optsVerbose ?options
      then let text = path ++ ": " ++ ansiGreen ++ "OK" ++ ansiReset ++ " " ++
                      prettyDuration False (realToFrac (diffUTCTime endtm starttm))
           in outputConcurrent (text ++ "\n")
      else return ()
    putMVar outvar success

  return ()

runTreePar topmparregion revidxlist revpath toptree@Resource{} topoutvar = runResource topmparregion 1 toptree topoutvar
  where
    runResource
      :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int, ?istty :: Bool, ?pool :: WorkerPool [Int])
      => Maybe ConsoleRegion -> Int -> TestTree -> MVar Bool -> IO ()
    runResource mparregion depth (Resource inhname make cleanup fun) outvar = do
      let pathitem = '[' : show depth ++ "](" ++ inhname ++ ")"
          path = intercalate "/" (reverse (pathitem : revpath))
          idxlist = reverse revidxlist
      -- outputConcurrent $ "! " ++ path ++ ": R Submitting\n"
      submitOrRunIn mparregion idxlist Nothing $ \makeRegion -> do
        setConsoleRegion makeRegion ('|' : path ++ " [R] making...")

        -- outputConcurrent $ "! " ++ path ++ ": R Making\n"
        value <- make  -- TODO: catch exceptions
        -- outputConcurrent $ "! " ++ path ++ ": R Made\n"

        -- outputConcurrent $ "! " ++ path ++ ": R Running subtree\n"
        suboutvar <- newEmptyMVar
        runResource (Just makeRegion) (depth + 1) (fun value) suboutvar  -- will consume makeRegion
        -- outputConcurrent $ "! " ++ path ++ ": R Scheduled subtree\n"

        _ <- forkIO $ do
          success <- readMVar suboutvar
          -- outputConcurrent $ "! " ++ path ++ ": R Subtree done, scheduling cleanup\n"
          poolSubmit ?pool idxlist (Just outvar) $ do
            cleanupRegion <- openConsoleRegion Linear
            setConsoleRegion cleanupRegion ('|' : path ++ " [R] cleanup...")
            -- outputConcurrent $ "! " ++ path ++ ": R Cleaning up\n"
            cleanup value  -- TODO: catch exceptions
            -- outputConcurrent $ "! " ++ path ++ ": R Cleanup done\n"
            closeConsoleRegion cleanupRegion
            return success
        return ()
    runResource mparregion _ tree outvar = runTreePar mparregion revidxlist revpath tree outvar

runTreePar mparregion revidxlist revpath (HP name prop) outvar = do
  let path = intercalate "/" (reverse (name : revpath))
      idxlist = reverse revidxlist

  submitOrRunIn mparregion idxlist (Just outvar) $ \region -> do
    -- outputConcurrent $ "! " ++ path ++ ": Running"
    let prefix = path ++ " [T]"
    setConsoleRegion region ('|' : prefix)
    let progressHandler report = do
          str <- H.renderProgress H.EnableColor (Just (fromString "")) report
          setConsoleRegion region ('|' : prefix ++ ": " ++ replace '\n' "  " str)
    (ok, rendered) <- runHP progressHandler revpath name prop
    finishConsoleRegion region (path ++ ": " ++ rendered)
    return ok

submitOrRunIn :: (?pool :: WorkerPool [Int])
              => Maybe ConsoleRegion -> [Int] -> Maybe (MVar a) -> (ConsoleRegion -> IO a) -> IO ()
submitOrRunIn Nothing idxlist outvar fun =
  poolSubmit ?pool idxlist outvar (openConsoleRegion Linear >>= fun)
submitOrRunIn (Just reg) _idxlist outvar fun = do
  result <- fun reg
  forM_ outvar $ \mvar -> putMVar mvar result

-- | If all tests are successful, returns the number of output lines produced
runTreeSeq :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int, ?istty :: Bool)
           => Int -> [String] -> TestTree -> IO (Maybe Int)
runTreeSeq indent revpath (Group opts name trees) = do
  putStrLn (replicate (2 * indent) ' ' ++ name) >> hFlush stdout
  starttm <- getCurrentTime
  mlns <- fmap (fmap sum . sequence) . forM trees $
            runTreeSeq (indent + 1) (name : revpath)
  endtm <- getCurrentTime
  case mlns of
    Just lns | goCollapse opts, ?istty -> do
      let thislen = 2*indent + length name
      putStrLn $ concat (replicate (lns+1) (ANSI.cursorUpCode 1 ++ ANSI.clearLineCode)) ++
                 ANSI.setCursorColumnCode 0 ++
                 replicate (2 * indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ' ++
                 ansiGreen ++ "OK" ++ ansiReset ++
                 prettyDuration False (realToFrac (diffUTCTime endtm starttm))
      return (Just 1)
    _ -> return ((+1) <$> mlns)
runTreeSeq indent path (Resource _ make cleanup fun) = do
  value <- make
  success <- runTreeSeq indent path (fun value)
  cleanup value
  return success
runTreeSeq indent path (HP name prop) = do
  let thislen = 2*indent + length name
  let prefix = replicate (2*indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' '
  when ?istty $ putStr prefix >> hFlush stdout
  (ok, rendered) <- runHP (outputProgress (?maxlen + 2)) path name prop
  putStrLn ((if ?istty then ANSI.clearFromCursorToLineEndCode else prefix) ++ rendered) >> hFlush stdout
  return (if ok then Just 1 else Nothing)

runHP :: (?options :: Options, ?stats :: IORef Stats, ?maxlen :: Int)
      => (H.Report H.Progress -> IO ())
      -> [String]
      -> String -> H.Property -> IO (Bool, String)
runHP progressPrinter revpath name (H.Property config test) = do
  let (config', seedfun) = applyHedgehogOptions ?options config
  seed <- seedfun

  starttm <- getCurrentTime
  report <- H.checkReport config' 0 seed test progressPrinter
  endtm <- getCurrentTime

  rendered <- renderResult report (intercalate "/" (reverse (name : revpath))) (diffUTCTime endtm starttm)

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

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 ++ ANSI.setCursorColumnCode indent)
      hFlush stdout
  | otherwise = return ()

renderResult :: 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 $ ansiGreen ++ "All " ++ show (statsTotal stats) ++
                 " tests passed." ++ prettyDuration True (realToFrac timeTaken) ++ ansiReset
  | otherwise =
      let nfailed = statsTotal stats - statsOK stats
      in putStrLn $ ansiRed ++ "Failed " ++ show nfailed ++ " out of " ++ show (statsTotal stats) ++
                    " tests." ++ prettyDuration True (realToFrac timeTaken) ++ ansiReset


newtype WorkerPool k = WorkerPool (TVar (PQ.MinPQueue k (Terminate PoolJob)))
data PoolJob = forall a. PoolJob (IO a) (Maybe (MVar a))
data Terminate a = Terminate | Value a
  deriving (Eq, Ord)  -- Terminate sorts before Value

withWorkerPool :: Ord k => Int -> (WorkerPool k -> IO a) -> IO a
withWorkerPool numWorkers k = do
  chan <- newTVarIO PQ.empty
  threads <- forM [0..numWorkers-1] (\i -> forkOn i (worker i chan))
  k (WorkerPool chan) `finally` do
    replicateM_ numWorkers (atomically $ writeTVar chan PQ.empty)
    forM_ threads killThread
  where
    worker :: Ord k => Int -> TVar (PQ.MinPQueue k (Terminate PoolJob)) -> IO ()
    worker idx chan = do
      job <- atomically $
        readTVar chan >>= \case
          PQ.Empty -> retry
          (_, j) PQ.:< q -> writeTVar chan q >> return j
      case job of
        Value (PoolJob action mmvar) -> do
          -- outputConcurrent $ "[" ++ show idx ++ "] got job\n"
          result <- action
          forM_ mmvar $ \mvar -> putMVar mvar result
          worker idx chan
        Terminate -> return ()

poolSubmit :: Ord k => WorkerPool k -> k -> Maybe (MVar a) -> IO a -> IO ()
poolSubmit (WorkerPool chan) key mmvar action =
  atomically $ modifyTVar chan $ PQ.insert key (Value (PoolJob action mmvar))


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 ""

ansiRed, ansiGreen, ansiReset :: (?istty :: Bool) => String
ansiRed   = ansi (ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull ANSI.Red])
ansiGreen = ansi (ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Dull ANSI.Green])
ansiReset = ansi (ANSI.setSGRCode [ANSI.Reset])

-- getTermIsDark :: IO Bool
-- getTermIsDark = do
--   mclr <- ANSI.getLayerColor ANSI.Background
--   case mclr of
--     Nothing -> return True
--     Just (RGB r g b) ->
--       let cvt n = fromIntegral n / fromIntegral (maxBound `asTypeOf` n)
--       in return $ (cvt r + cvt g + cvt b) / 3 < (0.5 :: Double)

-- ansiRegionBg :: (?istty :: Bool, ?termisdark :: Bool) => String
-- ansiRegionBg
--   | not ?istty = ""
--   | ?termisdark = ANSI.setSGRCode [ANSI.SetRGBColor ANSI.Background (rgb 0.0 0.05 0.1)]
--   | otherwise = ANSI.setSGRCode [ANSI.SetRGBColor ANSI.Background (rgb 0.95 0.95 1.0)]