{-# LANGUAGE OverloadedStrings, FlexibleContexts, LambdaCase #-}
-- | This program is a convenience utility for running the Futhark
-- test suite, and its test programs.
module Futhark.CLI.Test (main) where

import Control.Applicative.Lift (runErrors, failure, Errors, Lift(..))
import Control.Concurrent
import Control.Exception
import Control.Monad
import Control.Monad.Except hiding (throwError)
import qualified Control.Monad.Except as E
import qualified Data.ByteString as SBS
import qualified Data.ByteString.Lazy as LBS

import Data.List
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.IO as T
import System.Console.ANSI
import System.Process.ByteString (readProcessWithExitCode)
import System.Exit
import System.FilePath
import System.Console.GetOpt
import System.IO
import Text.Regex.TDFA

import Futhark.Analysis.Metrics
import Futhark.Test
import Futhark.Util.Options
import Futhark.Util.Pretty (prettyText)
import Futhark.Util.Table

--- Test execution

type TestM = ExceptT [T.Text] IO

-- Taken from transformers-0.5.5.0.
eitherToErrors :: Either e a -> Errors e a
eitherToErrors = either failure Pure

throwError :: MonadError [e] m => e -> m a
throwError e = E.throwError [e]

runTestM :: TestM () -> IO TestResult
runTestM = fmap (either Failure $ const Success) . runExceptT

io :: IO a -> TestM a
io = liftIO

context :: T.Text -> TestM a -> TestM a
context s = withExceptT $
  \case
    []      -> []
    (e:es') -> (s <> ":\n" <> e):es'

accErrors :: [TestM a] -> TestM [a]
accErrors tests = do
  eithers <- lift $ mapM runExceptT tests
  let errors = traverse eitherToErrors eithers
  ExceptT $ return $ runErrors errors

accErrors_ :: [TestM a] -> TestM ()
accErrors_ = void . accErrors

data TestResult = Success
                | Failure [T.Text]
                deriving (Eq, Show)

data TestCase = TestCase { _testCaseMode :: TestMode
                         , testCaseProgram :: FilePath
                         , testCaseTest :: ProgramTest
                         , _testCasePrograms :: ProgConfig
                         }
                deriving (Show)

instance Eq TestCase where
  x == y = testCaseProgram x == testCaseProgram y

instance Ord TestCase where
  x `compare` y = testCaseProgram x `compare` testCaseProgram y

data RunResult = ErrorResult Int SBS.ByteString
               | SuccessResult [Value]

progNotFound :: T.Text -> T.Text
progNotFound s = s <> ": command not found"

optimisedProgramMetrics :: ProgConfig -> StructurePipeline -> FilePath -> TestM AstMetrics
optimisedProgramMetrics programs pipeline program =
  case pipeline of SOACSPipeline ->
                     check "-s"
                   KernelsPipeline ->
                     check "--kernels"
                   SequentialCpuPipeline ->
                     check "--cpu"
                   GpuPipeline ->
                     check "--gpu"
  where check opt = do
          (code, output, err) <-
            io $ readProcessWithExitCode (configFuthark programs) ["dev", opt, "--metrics", program] ""
          let output' = T.decodeUtf8 output
          case code of
            ExitSuccess
              | [(m, [])] <- reads $ T.unpack output' -> return m
              | otherwise -> throwError $ "Could not read metrics output:\n" <> output'
            ExitFailure 127 -> throwError $ progNotFound $ T.pack $ configFuthark programs
            ExitFailure _ -> throwError $ T.decodeUtf8 err

testMetrics :: ProgConfig -> FilePath -> StructureTest -> TestM ()
testMetrics programs program (StructureTest pipeline (AstMetrics expected)) =
  context "Checking metrics" $ do
    actual <- optimisedProgramMetrics programs pipeline program
    accErrors_ $ map (ok actual) $ M.toList expected
  where ok (AstMetrics metrics) (name, expected_occurences) =
          case M.lookup name metrics of
            Nothing
              | expected_occurences > 0 ->
              throwError $ name <> " should have occurred " <> T.pack (show expected_occurences) <>
              " times, but did not occur at all in optimised program."
            Just actual_occurences
              | expected_occurences /= actual_occurences ->
                throwError $ name <> " should have occurred " <> T.pack (show expected_occurences) <>
              " times, but occured " <> T.pack (show actual_occurences) <> " times."
            _ -> return ()

testWarnings :: [WarningTest] -> SBS.ByteString -> TestM ()
testWarnings warnings futerr = accErrors_ $ map testWarning warnings
  where testWarning (ExpectedWarning regex_s regex)
          | not (match regex $ T.unpack $ T.decodeUtf8 futerr) =
            throwError $ "Expected warning:\n  " <> regex_s <>
            "\nGot warnings:\n  " <> T.decodeUtf8 futerr
          | otherwise = return ()

runTestCase :: TestCase -> TestM ()
runTestCase (TestCase mode program testcase progs) =
  case testAction testcase of

    CompileTimeFailure expected_error ->
      context (mconcat ["Type-checking with '", T.pack futhark,
                        " check ", T.pack program, "'"]) $ do
        (code, _, err) <-
          io $ readProcessWithExitCode futhark ["check", program] ""
        case code of
         ExitSuccess -> throwError "Expected failure\n"
         ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark
         ExitFailure 1 -> throwError $ T.decodeUtf8 err
         ExitFailure _ -> checkError expected_error err

    RunCases _ _ warnings | mode == TypeCheck -> do
      let options = ["check", program] ++ configExtraCompilerOptions progs
      context (mconcat ["Type-checking with '", T.pack futhark,
                        " check ", T.pack program, "'"]) $ do
        (code, _, err) <- io $ readProcessWithExitCode futhark options ""
        testWarnings warnings err
        case code of
         ExitSuccess -> return ()
         ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark
         ExitFailure _ -> throwError $ T.decodeUtf8 err

    RunCases ios structures warnings -> do
      -- Compile up-front and reuse same executable for several entry points.
      let backend = configBackend progs
          extra_options = configExtraCompilerOptions progs
      unless (mode == Compile) $
        context "Generating reference outputs" $
        ensureReferenceOutput futhark "c" program ios
      unless (mode == Interpreted) $
        context ("Compiling with --backend=" <> T.pack backend) $ do
          compileTestProgram extra_options futhark backend program warnings
          mapM_ (testMetrics progs program) structures
          unless (mode == Compile) $ do
            (tuning_opts, _) <-
              liftIO $ determineTuning (configTuning progs) program
            let progs' = progs { configExtraOptions =
                                 tuning_opts ++ configExtraOptions progs }
            context "Running compiled program" $
              accErrors_ $ map (runCompiledEntry program progs') ios
      unless (mode == Compile || mode == Compiled) $
        context "Interpreting" $
          accErrors_ $ map (runInterpretedEntry futhark program) ios
  where futhark = configFuthark progs

runInterpretedEntry :: String -> FilePath -> InputOutputs -> TestM()
runInterpretedEntry futhark program (InputOutputs entry run_cases) =
  let dir = takeDirectory program
      runInterpretedCase run@(TestRun _ inputValues _ index _) =
        unless ("compiled" `elem` runTags run) $
          context ("Entry point: " <> entry
                   <> "; dataset: " <> T.pack (runDescription run)) $ do

            input <- T.unlines . map prettyText <$> getValues dir inputValues
            expectedResult' <- getExpectedResult program entry run
            (code, output, err) <-
              io $ readProcessWithExitCode futhark ["run", "-e", T.unpack entry, program] $
              T.encodeUtf8 input
            case code of
              ExitFailure 127 -> throwError $ progNotFound $ T.pack futhark

              _               -> compareResult entry index program expectedResult'
                                 =<< runResult program code output err

  in accErrors_ $ map runInterpretedCase run_cases

runCompiledEntry :: FilePath -> ProgConfig -> InputOutputs -> TestM ()
runCompiledEntry program progs (InputOutputs entry run_cases) =
  -- Explicitly prefixing the current directory is necessary for
  -- readProcessWithExitCode to find the binary when binOutputf has
  -- no path component.
  let binOutputf = dropExtension program
      binpath = "." </> binOutputf
      entry_options = ["-e", T.unpack entry]

      runner = configRunner progs
      extra_options = configExtraOptions progs

      runCompiledCase run@(TestRun _ inputValues _ index _) =
        context ("Entry point: " <> entry
                 <> "; dataset: " <> T.pack (runDescription run)) $ do
          expected <- getExpectedResult program entry run
          (progCode, output, progerr) <-
            runProgram runner extra_options program entry inputValues
          compareResult entry index program expected
            =<< runResult program progCode output progerr

  in context ("Running " <> T.pack (unwords $ binpath : entry_options ++ extra_options)) $
     accErrors_ $ map runCompiledCase run_cases

checkError :: ExpectedError -> SBS.ByteString -> TestM ()
checkError (ThisError regex_s regex) err
  | not (match regex $ T.unpack $ T.decodeUtf8 err) =
     throwError $ "Expected error:\n  " <> regex_s <>
     "\nGot error:\n  " <> T.decodeUtf8 err
checkError _ _ =
  return ()

runResult :: FilePath -> ExitCode -> SBS.ByteString -> SBS.ByteString -> TestM RunResult
runResult program ExitSuccess stdout_s _ =
  case valuesFromByteString "stdout" $ LBS.fromStrict stdout_s of
    Left e   -> do
      let actualf = program `addExtension` "actual"
      io $ SBS.writeFile actualf stdout_s
      throwError $ T.pack e <> "\n(See " <> T.pack actualf <> ")"
    Right vs -> return $ SuccessResult vs
runResult _ (ExitFailure code) _ stderr_s =
  return $ ErrorResult code stderr_s

compileTestProgram :: [String] -> FilePath -> String -> FilePath -> [WarningTest] -> TestM ()
compileTestProgram extra_options futhark backend program warnings = do
  (_, futerr) <- compileProgram extra_options futhark backend program
  testWarnings warnings futerr

compareResult :: T.Text -> Int -> FilePath -> ExpectedResult [Value] -> RunResult
              -> TestM ()
compareResult _ _ _ (Succeeds Nothing) SuccessResult{} =
  return ()
compareResult entry index program (Succeeds (Just expectedResult)) (SuccessResult actualResult) =
  case compareValues1 actualResult expectedResult of
    Just mismatch -> do
      let actualf = program <.> T.unpack entry <.> show index <.> "actual"
          expectedf = program <.> T.unpack entry <.> show index <.> "expected"
      io $ SBS.writeFile actualf $
        T.encodeUtf8 $ T.unlines $ map prettyText actualResult
      io $ SBS.writeFile expectedf $
        T.encodeUtf8 $ T.unlines $ map prettyText expectedResult
      throwError $ T.pack actualf <> " and " <> T.pack expectedf <>
        " do not match:\n" <> T.pack (show mismatch) <> "\n"
    Nothing ->
      return ()
compareResult _ _ _ (RunTimeFailure expectedError) (ErrorResult _ actualError) =
  checkError expectedError actualError
compareResult _ _ _ (Succeeds _) (ErrorResult code err) =
  throwError $ "Program failed with error code " <>
  T.pack (show code) <> " and stderr:\n  " <> T.decodeUtf8 err
compareResult _ _ _ (RunTimeFailure f) (SuccessResult _) =
  throwError $ "Program succeeded, but expected failure:\n  " <> T.pack (show f)

---
--- Test manager
---

data TestStatus = TestStatus { testStatusRemain :: [TestCase]
                             , testStatusRun :: [TestCase]
                             , testStatusTotal :: Int
                             , testStatusFail :: Int
                             , testStatusPass :: Int
                             , testStatusRuns :: Int
                             , testStatusRunsRemain :: Int
                             , testStatusRunPass :: Int
                             , testStatusRunFail :: Int
                             }

catching :: IO TestResult -> IO TestResult
catching m = m `catch` save
  where save :: SomeException -> IO TestResult
        save e = return $ Failure [T.pack $ show e]

doTest :: TestCase -> IO TestResult
doTest = catching . runTestM . runTestCase

makeTestCase :: TestConfig -> TestMode -> (FilePath, ProgramTest) -> TestCase
makeTestCase config mode (file, spec) =
  TestCase mode file spec $ configPrograms config

data ReportMsg = TestStarted TestCase
               | TestDone TestCase TestResult

runTest :: MVar TestCase -> MVar ReportMsg -> IO ()
runTest testmvar resmvar = forever $ do
  test <- takeMVar testmvar
  putMVar resmvar $ TestStarted test
  res <- doTest test
  putMVar resmvar $ TestDone test res

excludedTest :: TestConfig -> TestCase -> Bool
excludedTest config =
  any (`elem` configExclude config) . testTags . testCaseTest

statusTable :: TestStatus -> String
statusTable ts = buildTable rows 1
  where rows =
          [ [ mkEntry "", passed, failed, mkEntry "remaining"]
          , map mkEntry ["programs", passedProgs, failedProgs, remainProgs']
          , map mkEntry ["runs", passedRuns, failedRuns, remainRuns']
          ]
        passed       = ("passed", [SetColor Foreground Vivid Green])
        failed       = ("failed", [SetColor Foreground Vivid Red])
        passedProgs  = show $ testStatusPass ts
        failedProgs  = show $ testStatusFail ts
        totalProgs   = show $ testStatusTotal ts
        totalRuns    = show $ testStatusRuns ts
        passedRuns   = show $ testStatusRunPass ts
        failedRuns   = show $ testStatusRunFail ts
        remainProgs  = show . length $ testStatusRemain ts
        remainProgs' = remainProgs ++ "/" ++ totalProgs
        remainRuns   = show $ testStatusRunsRemain ts
        remainRuns'  = remainRuns ++ "/" ++ totalRuns

tableLines :: Int
tableLines = 1 + (length . lines $ blankTable)
  where blankTable = statusTable $ TestStatus [] [] 0 0 0 0 0 0 0

spaceTable :: IO ()
spaceTable = putStr $ replicate tableLines '\n'

reportTable :: TestStatus -> IO ()
reportTable ts = do
  moveCursorToTableTop
  putStrLn $ statusTable ts
  clearLine
  putStrLn $ atMostChars 60 running
  where running    = "Now testing: " ++
                     (unwords . reverse . map testCaseProgram . testStatusRun) ts

moveCursorToTableTop :: IO ()
moveCursorToTableTop = cursorUpLine tableLines

atMostChars :: Int -> String -> String
atMostChars n s | length s > n = take (n-3) s ++ "..."
                | otherwise    = s

reportText :: TestStatus -> IO ()
reportText ts =
  putStr $ "(" ++ show (testStatusFail ts)  ++ " failed, " ++
                  show (testStatusPass ts)  ++ " passed, " ++
                  show num_remain           ++ " to go).\n"
    where num_remain  = length $ testStatusRemain ts

runTests :: TestConfig -> [FilePath] -> IO ()
runTests config paths = do
  -- We force line buffering to ensure that we produce running output.
  -- Otherwise, CI tools and the like may believe we are hung and kill
  -- us.
  hSetBuffering stdout LineBuffering

  let mode = configTestMode config
  all_tests <- map (makeTestCase config mode) <$> testSpecsFromPaths paths
  testmvar <- newEmptyMVar
  reportmvar <- newEmptyMVar
  concurrency <- getNumCapabilities
  replicateM_ concurrency $ forkIO $ runTest testmvar reportmvar

  let (excluded, included) = partition (excludedTest config) all_tests
  _ <- forkIO $ mapM_ (putMVar testmvar) included
  isTTY <- (&& not (configLineOutput config)) <$> hIsTerminalDevice stdout

  let report | isTTY = reportTable
             | otherwise = reportText
      clear | isTTY = clearFromCursorToScreenEnd
            |otherwise = putStr "\n"

      numTestCases tc =
        case testAction $ testCaseTest tc of
          CompileTimeFailure _ -> 1
          RunCases ios sts wts -> (length . concat) (iosTestRuns <$> ios)
                                  + length sts + length wts

      getResults ts
        | null (testStatusRemain ts) = report ts >> return ts
        | otherwise = do
          report ts
          msg <- takeMVar reportmvar
          case msg of
            TestStarted test -> do
              unless isTTY $
                putStr $ "Started testing " <> testCaseProgram test <> " "
              getResults $ ts {testStatusRun = test : testStatusRun ts}
            TestDone test res -> do
              let ts' = ts { testStatusRemain = test `delete` testStatusRemain ts
                           , testStatusRun    = test `delete` testStatusRun ts
                           , testStatusRunsRemain = testStatusRunsRemain ts
                                                    - numTestCases test
                           }
              case res of
                Success -> do
                  let ts'' = ts' { testStatusRunPass =
                                     testStatusRunPass ts' + numTestCases test
                                 }
                  unless isTTY $
                    putStr $ "Finished testing " <> testCaseProgram test <> " "
                  getResults $ ts'' { testStatusPass = testStatusPass ts + 1}
                Failure s -> do
                  when isTTY moveCursorToTableTop
                  clear
                  T.putStr $ (T.pack (inRed $ testCaseProgram test) <> ":\n") <> T.unlines s
                  when isTTY spaceTable
                  getResults $ ts' { testStatusFail = testStatusFail ts' + 1
                                   , testStatusRunPass = testStatusRunPass ts'
                                                         + numTestCases test - length s

                                   , testStatusRunFail = testStatusRunFail ts'
                                                         + length s
                                   }

  when isTTY spaceTable

  ts <- getResults TestStatus { testStatusRemain = included
                              , testStatusRun    = []
                              , testStatusTotal  = length included
                              , testStatusFail   = 0
                              , testStatusPass   = 0
                              , testStatusRuns  = sum $ map numTestCases included
                              , testStatusRunsRemain = sum $ map numTestCases included
                              , testStatusRunPass = 0
                              , testStatusRunFail = 0
                              }

  -- Removes "Now testing" output.
  when isTTY $ cursorUpLine 1 >> clearLine

  let excluded_str | null excluded = ""
                   | otherwise = " (" ++ show (length excluded) ++ " program(s) excluded).\n"
  putStr excluded_str
  exitWith $ case testStatusFail ts of 0 -> ExitSuccess
                                       _ -> ExitFailure 1

inRed :: String -> String
inRed s = setSGRCode [SetColor Foreground Vivid Red] ++ s ++ setSGRCode [Reset]

---
--- Configuration and command line parsing
---

data TestConfig = TestConfig
                  { configTestMode :: TestMode
                  , configPrograms :: ProgConfig
                  , configExclude :: [T.Text]
                  , configLineOutput :: Bool
                  }

defaultConfig :: TestConfig
defaultConfig = TestConfig { configTestMode = Everything
                           , configExclude = [ "disable" ]
                           , configPrograms =
                             ProgConfig
                             { configBackend = "c"
                             , configFuthark = "futhark"
                             , configRunner = ""
                             , configExtraOptions = []
                             , configExtraCompilerOptions = []
                             , configTuning = Just "tuning"
                             }
                           , configLineOutput = False
                           }

data ProgConfig = ProgConfig
                  { configBackend :: String
                  , configFuthark :: FilePath
                  , configRunner :: FilePath
                  , configExtraCompilerOptions :: [String]
                  , configTuning :: Maybe String
                  , configExtraOptions :: [String]
                  -- ^ Extra options passed to the programs being run.
                  }
                  deriving (Show)

changeProgConfig :: (ProgConfig -> ProgConfig) -> TestConfig -> TestConfig
changeProgConfig f config = config { configPrograms = f $ configPrograms config }

setBackend :: FilePath -> ProgConfig -> ProgConfig
setBackend backend config =
  config { configBackend = backend }

setFuthark :: FilePath -> ProgConfig -> ProgConfig
setFuthark futhark config =
  config { configFuthark = futhark }

setRunner :: FilePath -> ProgConfig -> ProgConfig
setRunner runner config =
  config { configRunner = runner }

addCompilerOption :: String -> ProgConfig -> ProgConfig
addCompilerOption option config =
  config { configExtraCompilerOptions = configExtraCompilerOptions config ++ [option] }

addOption :: String -> ProgConfig -> ProgConfig
addOption option config =
  config { configExtraOptions = configExtraOptions config ++ [option] }

data TestMode = TypeCheck
              | Compile
              | Compiled
              | Interpreted
              | Everything
              deriving (Eq, Show)

commandLineOptions :: [FunOptDescr TestConfig]
commandLineOptions = [
    Option "t" ["typecheck"]
    (NoArg $ Right $ \config -> config { configTestMode = TypeCheck })
    "Only perform type-checking"
  , Option "i" ["interpreted"]
    (NoArg $ Right $ \config -> config { configTestMode = Interpreted })
    "Only interpret"
  , Option "c" ["compiled"]
    (NoArg $ Right $ \config -> config { configTestMode = Compiled })
    "Only run compiled code"
  , Option "C" ["compile"]
    (NoArg $ Right $ \config -> config { configTestMode = Compile })
    "Only compile, do not run."
  , Option [] ["no-terminal", "notty"]
    (NoArg $ Right $ \config -> config { configLineOutput = True })
    "Provide simpler line-based output."
  , Option [] ["backend"]
    (ReqArg (Right . changeProgConfig . setBackend) "BACKEND")
    "Backend used for compilation (defaults to 'c')."
  , Option [] ["futhark"]
    (ReqArg (Right . changeProgConfig . setFuthark) "PROGRAM")
    "Program to run for subcommands (defaults to 'futhark')."
  , Option [] ["runner"]
    (ReqArg (Right . changeProgConfig . setRunner) "PROGRAM")
    "The program used to run the Futhark-generated programs (defaults to nothing)."
  , Option [] ["exclude"]
    (ReqArg (\tag ->
               Right $ \config ->
               config { configExclude = T.pack tag : configExclude config })
     "TAG")
    "Exclude test programs that define this tag."
  , Option "p" ["pass-option"]
    (ReqArg (Right . changeProgConfig . addOption) "OPT")
    "Pass this option to programs being run."
  , Option [] ["pass-compiler-option"]
    (ReqArg (Right . changeProgConfig . addCompilerOption) "OPT")
    "Pass this option to the compiler (or typechecker if in -t mode)."
  ]

main :: String -> [String] -> IO ()
main = mainWithOptions defaultConfig commandLineOptions "options... programs..." $ \progs config ->
  Just $ runTests config progs