module Benchmark where

import Definitions
import Shellish hiding ( run )
import Data.Char
import Data.List
import Data.List.Split ( wordsBy )
import System.Directory
import System.Environment
import System.FilePath( (</>), (<.>), splitDirectories, joinPath )
import System.IO
import System.Exit
import Text.Regex.Posix( (=~) )
import Data.Time.Clock
import Control.Monad.Error
import Control.Exception( throw )
import Control.Concurrent( forkIO )
import Control.Concurrent.Chan( newChan, writeChan, readChan, Chan )
import System.Console.CmdArgs (isLoud)
import System.Process( runInteractiveProcess, runInteractiveCommand,
                       waitForProcess )
import qualified System.IO.UTF8 as UTF8

copyTree :: FilePath -> FilePath -> IO ()
copyTree from to =
    do subs <- (\\ [".", ".."]) `fmap` getDirectoryContents from
       createDirectory to
       forM_ subs $ \item -> do
         is_dir <- doesDirectoryExist (from </> item)
         is_file <- doesFileExist (from </> item)
         when is_dir $ copyTree (from </> item) (to </> item)
         when is_file $ copyFile (from </> item) (to </> item)

reset :: Command ()
reset = do
  resetMemoryUsed
  resetTimeUsed

exec :: Benchmark a -> FilePath -> TestRepo -> Command a
exec (Idempotent _ _ cmd) darcs_path tr = do
  cd "_playground"
  verbose "cd _playground"
  cmd (darcs darcs_path) tr
exec (Destructive _ _ cmd) darcs_path tr = do
  cd "_playground"
  let cleanup = verbose "cd .. ; rm -rf _playground" >> cd ".." >> rm_rf "_playground"
  res <- cmd (darcs darcs_path) tr `catchError` \e -> (cleanup >> throw e)
  cleanup
  return res
exec (Description _) _ _ = fail "Cannot run description-only benchmark."

defaultrepo, sources :: FilePath -> FilePath
defaultrepo path = path </> "_darcs" </> "prefs" </> "defaultrepo"
sources path = path </> "_darcs" </> "prefs" </> "sources"

prepare :: String -> Command ()
prepare origrepo = do
  progress "!" >> verbose "rm -rf _playground"
  rm_rf "_playground"
  liftIO $ createDirectory "_playground"
  let playrepo = "_playground" </> "repo"
  isrepo <- liftIO $ doesDirectoryExist (origrepo </> "_darcs")
  unless isrepo $ fail $ origrepo ++ ": Not a darcs repository!"
  progress "." >> verbose ("cp -a '" ++ origrepo ++ "' '" ++ playrepo ++ "'")
  liftIO $ copyTree origrepo playrepo
  progress "." >> verbose ("# sanitize " ++ playrepo)
  wd <- pwd
  liftIO $ do writeFile (defaultrepo playrepo) (wd </> origrepo)
              removeFile (sources playrepo) `catch` \_ -> return ()

prepareIfDifferent :: String -> Command ()
prepareIfDifferent origrepo = do
  let playrepo = "_playground" </> "repo"
  exist <- test_e "_playground"
  current' <- if exist then liftIO $ readFile (defaultrepo playrepo) else return ""
  let current = reverse (dropWhile (=='\n') $ reverse current')
  wd <- pwd
  if (exist && current == wd </> origrepo)
     then progress "..." >> verbose ("# leaving " ++ playrepo ++ " alone")
     else prepare origrepo

-- | Run a benchmark as many times as it takes to pass a minimum threshold
--   of time or iterations (whichever comes first)
--   Useful for very small benchmarks
adaptive :: Double -- ^ seconds
         -> (Int,Int) -- ^ min, max iterations
         -> Command MemTime -> Command [MemTime]
adaptive thresh (iters_min,iters_max) cmd =
   cmd >> -- just once to warm up the disk cache (eliminate a source of variance)
   go thresh iters_max []
 where
  go t i acc | i <= 0 || (t <= 0 && i <= iters_enough) = return acc
  go t i acc =
   do verbose $ "# adaptive: iterations remaining: " ++ show i ++ " time remaining: " ++ show t
      mt@(MemTime _ t2) <- cmd
      go (t - t2) (i - 1) (mt : acc)
  iters_enough = iters_max - iters_min

run :: Test a -> Command (Maybe MemTimeOutput)
run test@(Test benchmark tr (TestBinary bin)) = do
  (Just `fmap` run') `catchError` \e ->
      do echo_n_err $ " error: " ++ show e
         return Nothing
  where run' = do
          progress $ bin ++ " " ++ description benchmark ++ " [" ++ trName tr ++ "]: "
          verbose $ "\n# testing; binary = " ++ bin ++ ", benchmark = " ++
                    description benchmark ++ ", repository = " ++ trName tr
          exe <- which $ bin
          darcs_path <- case exe of
                          Nothing -> canonize bin
                          Just p -> return p
          times <- adaptive 10 (3,100) . sub $ do
                     prepareIfDifferent (trPath tr)
                     liftIO (maybeVMFlush darcs_path)
                     m <- timed (exec benchmark darcs_path tr)
                     return m
          let result = mkMemTimeOutput times
              spaces = 45 - (length bin + length (description benchmark) + length (trName tr))
              tu = appropriateUnit (mtTimeMean result)
              result_str = unwords $ concatMap (\f -> f tu (Cell result)) [ formatTimeResult, formatMemoryResult, formatSampleSize ]
          liftIO $ appendResult test times
          progress $ (replicate spaces ' ') ++ result_str ++ "\n"
          verbose $ "# result: " ++ result_str
          return result

timed :: Command a -> Command MemTime
timed a = do
  resetMemoryUsed
  t1 <- liftIO $ getCurrentTime
  _ <- a
  t2 <- liftIO $ getCurrentTime
  mem <- memoryUsed
  resetMemoryUsed
  return $ MemTime (fromIntegral mem) (realToFrac $ diffUTCTime t2 t1)

darcsVersion :: String -> IO Version
darcsVersion cmd = do
       (_,outH,_,procH) <- runInteractiveCommand $ cmd ++ " --version"
       out <- strictGetContents outH
       _ <- waitForProcess procH
       return $ map read . wordsBy (== '.') . takeWhile (not . isSpace) $ out

check_darcs :: String -> IO ()
check_darcs cmd = do
       out <- darcsVersion cmd
       case out of
         2:_ -> return ()
         _ -> fail $ cmd ++ ": Not darcs 2.x binary."

verbose :: String -> Command ()
verbose m = liftIO $ do loud <- isLoud
                        when loud $ UTF8.hPutStrLn stderr m

progress :: String -> Command ()
progress m = liftIO $ do loud <- isLoud
                         unless loud $ UTF8.hPutStr stderr m

drain :: Handle -> Bool -> IO (Chan String)
drain h verb = do chan <- newChan
                  let work acc = do line <- hGetLine h
                                    when verb $ putStrLn ("## " ++ line)
                                    work (acc ++ line)
                              `catch` \_ -> writeChan chan acc
                  _ <- forkIO $ work ""
                  return chan

readFile' :: FilePath -> IO String
readFile' f = do s <- readFile f
                 length s `seq` return s

darcs :: String -> [String] -> Command String
darcs cmd args' = do
    stats_f <- liftIO $
      do tmpdir <- getTemporaryDirectory
         (f, h) <- openTempFile tmpdir "darcs-stats-XXXX"
         hClose h
         return f
    let args = args' ++ ["+RTS", "-s" ++ stats_f, "-RTS"]
    loud <- liftIO isLoud
    verbose . unwords $ cmd:args
    (res, _, stats) <- liftIO $ do
       mPlayground <- seekPlayground `fmap` getCurrentDirectory
       mEnv <- case mPlayground of
                 Nothing -> return Nothing
                 Just p  -> (Just . replace "HOME" p) `fmap` getEnvironment
       (_,outH,errH,procH) <- runInteractiveProcess cmd args Nothing mEnv
       res' <- drain outH loud
       errs' <- drain errH loud
       ex <- waitForProcess procH
       stats <- do c <- readFile' stats_f
                   removeFile stats_f `catch` \e -> hPutStrLn stderr (show e)
                   return c
                `catch` \_ -> return ""
       errs <- readChan errs'
       case ex of
         ExitSuccess -> return ()
         ExitFailure n -> fail $ "darcs failed with error code "
                            ++ show n ++ "\nsaying: " ++ errs
       res <- readChan res'
       return (res, errs, stats)
    let bytes = (stats =~ "([0-9, ]+) M[bB] total memory in use") :: String
        mem = (read (filter (`elem` "0123456789") bytes) :: Int)
    recordMemoryUsed $ mem * 1024 * 1024
    return res
 where
  replace k v xs = (k,v) : filter ((/= k) . fst) xs

seekPlayground :: FilePath -> Maybe FilePath
seekPlayground dir =
 if playground `elem` pieces
    then Just . joinPath . reverse . dropWhile (/= playground) . reverse $ pieces
    else Nothing
 where
  pieces = splitDirectories dir
  playground = "_playground"

-- ----------------------------------------------------------------------
-- variants
-- ----------------------------------------------------------------------

mkVariant :: String -> String -> Variant -> Command ()
mkVariant origrepo darcs_path v =
 case vId v of
  OptimizePristineVariant -> do
    isrepo <- liftIO $ doesDirectoryExist (origrepo </> "_darcs")
    unless isrepo $ fail $ origrepo ++ ": Not a darcs repository!"
    variant_isrepo  <- liftIO $ doesDirectoryExist (variant_repo </> "_darcs")
    unless variant_isrepo $ do
      echo $ "Setting up " ++ vDescription v ++ " variant of " ++ origrepo
      verbose ("cp -a '" ++ origrepo ++ "' '" ++ variant_repo ++ "'")
      liftIO $ copyTree origrepo variant_repo
      verbose ("# sanitize " ++ variant_repo)
      liftIO $ removeFile (sources variant_repo) `catch` \_ -> return ()
      darcs darcs_path [ "optimize", "--pristine", "--repodir", variant_repo ]
      return ()
  DefaultVariant -> return ()
 where
  variant_repo = variantRepoName v origrepo

variantRepoName :: Variant -> String -> String
variantRepoName (Variant { vId = DefaultVariant }) x = x
variantRepoName v x = "variant" <.> stripped x ++ "-" ++ vSuffix v
 where
  stripped y | "repo."   `isPrefixOf` y = stripped (drop 5 y)
             | "-hashed" `isSuffixOf` y = take (length y -  7) y
             | otherwise = y

setupVariants :: [TestRepo] -> TestBinary -> Command ()
setupVariants repos (TestBinary bin) =
  sequence_ [ mkVariant (trPath repo) bin variant
            | repo <- repos, variant <- trVariants repo ]