module Distribution.Simple.Bench
    ( bench
    ) where
import qualified Distribution.PackageDescription as PD
    ( PackageDescription(..), BuildInfo(buildable)
    , Benchmark(..), BenchmarkInterface(..), benchmarkType, hasBenchmarks )
import Distribution.Simple.BuildPaths ( exeExtension )
import Distribution.Simple.Compiler ( Compiler(..) )
import Distribution.Simple.InstallDirs
    ( fromPathTemplate, initialPathTemplateEnv, PathTemplateVariable(..)
    , substPathTemplate , toPathTemplate, PathTemplate )
import qualified Distribution.Simple.LocalBuildInfo as LBI
    ( LocalBuildInfo(..) )
import Distribution.Simple.Setup ( BenchmarkFlags(..), fromFlag )
import Distribution.Simple.UserHooks ( Args )
import Distribution.Simple.Utils ( die, notice, rawSystemExitCode )
import Distribution.Text
import Control.Monad ( when, unless )
import System.Exit ( ExitCode(..), exitFailure, exitWith )
import System.Directory ( doesFileExist )
import System.FilePath ( (</>), (<.>) )
bench :: Args                    
      -> PD.PackageDescription   
      -> LBI.LocalBuildInfo      
      -> BenchmarkFlags          
      -> IO ()
bench args pkg_descr lbi flags = do
    let verbosity         = fromFlag $ benchmarkVerbosity flags
        benchmarkNames    = args
        pkgBenchmarks     = PD.benchmarks pkg_descr
        enabledBenchmarks = [ t | t <- pkgBenchmarks
                            , PD.benchmarkEnabled t
                            , PD.buildable (PD.benchmarkBuildInfo t) ]
        
        doBench :: PD.Benchmark -> IO ExitCode
        doBench bm =
            case PD.benchmarkInterface bm of
              PD.BenchmarkExeV10 _ _ -> do
                  let cmd = LBI.buildDir lbi </> PD.benchmarkName bm
                            </> PD.benchmarkName bm <.> exeExtension
                      options = map (benchOption pkg_descr lbi bm) $
                                benchmarkOptions flags
                      name = PD.benchmarkName bm
                  
                  exists <- doesFileExist cmd
                  unless exists $ die $
                      "Error: Could not find benchmark program \""
                      ++ cmd ++ "\". Did you build the package first?"
                  notice verbosity $ startMessage name
                  
                  
                  exitcode <- rawSystemExitCode verbosity cmd options
                  notice verbosity $ finishMessage name exitcode
                  return exitcode
              _ -> do
                  notice verbosity $ "No support for running "
                      ++ "benchmark " ++ PD.benchmarkName bm ++ " of type: "
                      ++ show (disp $ PD.benchmarkType bm)
                  exitFailure
    when (not $ PD.hasBenchmarks pkg_descr) $ do
        notice verbosity "Package has no benchmarks."
        exitWith ExitSuccess
    when (PD.hasBenchmarks pkg_descr && null enabledBenchmarks) $
        die $ "No benchmarks enabled. Did you remember to configure with "
              ++ "\'--enable-benchmarks\'?"
    bmsToRun <- case benchmarkNames of
            [] -> return enabledBenchmarks
            names -> flip mapM names $ \bmName ->
                let benchmarkMap = zip enabledNames enabledBenchmarks
                    enabledNames = map PD.benchmarkName enabledBenchmarks
                    allNames = map PD.benchmarkName pkgBenchmarks
                in case lookup bmName benchmarkMap of
                    Just t -> return t
                    _ | bmName `elem` allNames ->
                          die $ "Package configured with benchmark "
                                ++ bmName ++ " disabled."
                      | otherwise -> die $ "no such benchmark: " ++ bmName
    let totalBenchmarks = length bmsToRun
    notice verbosity $ "Running " ++ show totalBenchmarks ++ " benchmarks..."
    exitcodes <- mapM doBench bmsToRun
    let allOk = totalBenchmarks == length (filter (== ExitSuccess) exitcodes)
    unless allOk exitFailure
  where
    startMessage name = "Benchmark " ++ name ++ ": RUNNING...\n"
    finishMessage name exitcode = "Benchmark " ++ name ++ ": "
                               ++ (case exitcode of
                                        ExitSuccess -> "FINISH"
                                        ExitFailure _ -> "ERROR")
benchOption :: PD.PackageDescription
            -> LBI.LocalBuildInfo
            -> PD.Benchmark
            -> PathTemplate
            -> String
benchOption pkg_descr lbi bm template =
    fromPathTemplate $ substPathTemplate env template
  where
    env = initialPathTemplateEnv
          (PD.package pkg_descr) (compilerId $ LBI.compiler lbi)
          (LBI.hostPlatform lbi) ++
          [(BenchmarkNameVar, toPathTemplate $ PD.benchmarkName bm)]