{-# LANGUAGE DeriveDataTypeable, ScopedTypeVariables #-}
module Benchmark where

import Prelude hiding ( readFile, catch )
import Definitions
import Shellish hiding ( run )
import Data.Typeable( Typeable )
import Data.Char
import Data.List
import Data.Maybe
import Data.DateTime( parseDateTime, startOfTime )
import Data.List.Split ( splitOn )
import System.Directory
import System.Environment
import System.FilePath( (</>), (<.>), splitDirectories, joinPath )
import System.IO hiding ( readFile )
import System.IO.Strict( readFile )
import System.Exit
import Text.Regex.Posix( (=~) )
import Data.Time.Clock
import Control.Monad.Error
import Control.Exception( Exception(..), throw, catch, SomeException )
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
import qualified Data.ByteString.Char8 as BS

catchany :: IO a -> (SomeException -> IO a) -> IO a
catchany = catch

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 -> TestBinary -> TestRepo -> Command a
exec (Idempotent _ _ cmd) bin tr = do
  cd "_playground"
  verbose "cd _playground"
  cmd (vcs bin) tr
exec (Destructive _ _ cmd) bin tr = do
  cd "_playground"
  let cleanup = verbose "cd .. ; rm -rf _playground" >> cd ".." >> rm_rf "_playground"
  res <- cmd (vcs bin) tr `catch_sh` \(e :: SomeException) -> (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")
  progress "." >> verbose ("cp -a '" ++ origrepo ++ "' '" ++ playrepo ++ "'")
  liftIO $ copyTree origrepo playrepo
  progress "." >> verbose ("# sanitize " ++ playrepo)
  wd <- pwd
  mkdir_p (playrepo </> "_darcs" </> "prefs") -- FIXME ... this is not very nice in git repos
  liftIO $ do writeFile (defaultrepo playrepo) (wd </> origrepo)
              removeFile (sources playrepo) `catchany` \_ -> return ()

prepareIfDifferent :: String -> Command ()
prepareIfDifferent origrepo = do
  let playrepo = "_playground" </> "repo"
  exist <- test_e "_playground"
  isrepo <- test_d $ playrepo </> "_darcs"
  if isrepo then do
    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
   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 bin) = do
  (Just `fmap` run') `catch_sh` \(e :: SomeException) ->
      do echo_err $ " error: " ++ show e
         return Nothing
  where run' = do
          progress $ cmd ++ " " ++ description benchmark ++ " [" ++ trName tr ++ "]: "
          verbose $ "\n# testing; binary = " ++ cmd ++ ", benchmark = " ++
                    description benchmark ++ ", repository = " ++ trName tr
          times <- adaptive 10 (3,100) . sub $ do
                     prepareIfDifferent (trPath tr)
                     binPath bin >>= liftIO . maybeVMFlush
                     m <- timed (exec benchmark bin tr)
                     return m
          let result = mkMemTimeOutput times
              spaces = 45 - (length cmd + 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
        cmd = binCommand bin

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)

darcsMeta :: String -> [String] -> IO String
darcsMeta cmd args = do
       (_,outH,_,procH) <- runInteractiveProcess cmd args Nothing Nothing
       out <- strictGetContents outH
       _ <- waitForProcess procH
       return out

check_vcs :: String -> IO TestBinary
check_vcs cmd = do
  ver <- darcsMeta cmd ["--version"]
  case ver of
    _ | "git" `isPrefixOf` ver ->
      return TestBinary { binCommand = cmd
                        , binVCS = VCSGit
                        , binVersionString = numeric ver
                        , binDate = startOfTime
                        , binGHC = "none at all"
                        , binContext = BS.empty }
    _ | "Mercurial" `isPrefixOf` ver ->
      return TestBinary { binCommand = cmd
                        , binVCS = VCSHg
                        , binVersionString = numeric ver
                        , binDate = startOfTime
                        , binGHC = "none at all"
                        , binContext = BS.empty }
    _ -> check_darcs cmd
  where numeric str = takeWhile (`elem` "1234567890.-") $ dropWhile (`notElem` "1234567890") str


check_darcs :: String -> IO TestBinary
check_darcs cmd = do
       version <- darcsMeta cmd ["--version"]
       [info, context] <- splitOn "Context:\n\n" `fmap` darcsMeta cmd ["--exact-version"]
       rts <- read `fmap` darcsMeta cmd ["+RTS", "--info"]
       let date' = case info of
             _ | "darcs compiled on" `isPrefixOf` info ->
               drop 18 . takeWhile (/='\n') $ info
             _ -> "<unknown date>"
           date = fromMaybe startOfTime $ parseDateTime "%b %e %Y, at %H:%M:%S" date'
           bin = TestBinary { binCommand = cmd
                            , binVCS = VCSDarcs
                            , binVersionString = takeWhile (/='\n') version
                            , binDate = date
                            , binGHC = fromMaybe "unknown" $ lookup "GHC version" rts
                            , binContext = BS.pack context }
       case binVersion bin of
         2:_ -> return bin
         _ -> 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 ++ "\n")
                              `catchany` \_ -> writeChan chan acc
                  _ <- forkIO $ work ""
                  return chan

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

data RunFailed = RunFailed String Int String deriving (Typeable)

instance Show RunFailed where
  show (RunFailed cmd code errs) =
    "error running " ++ cmd ++ ": exit status " ++ show code ++ ":\n" ++ errs

instance Exception RunFailed

runInPlayground :: String -> [String] -> Command String
runInPlayground cmd args = do
    loud <- liftIO isLoud
    verbose . unwords $ cmd:args
    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
       errs <- readChan errs'
       case ex of
         ExitSuccess -> return ()
         ExitFailure n -> throw $ RunFailed (cmd ++ " " ++ show args) n errs
       res <- readChan res'
       return res
 where
  replace k v xs = (k,v) : filter ((/= k) . fst) xs

binPath :: TestBinary -> Command String
binPath bin = do exe <- which (binCommand bin)
                 case exe of
                   Nothing -> canonize $ binCommand bin
                   Just p -> return p

vcs :: TestBinary -> [String] -> Command String
vcs bin args =
  do exe <- binPath bin
     case binVCS bin of
       VCSDarcs -> darcs exe args
       VCSGit -> git exe args
       VCSHg -> hg exe args

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"]
    res <- runInPlayground cmd args
    stats <- liftIO $ do c <- readFile' stats_f
                         removeFile stats_f `catchany` \e -> hPutStrLn stderr (show e)
                         return c
                       `catchany` \_ -> return ""
    let bytes = (stats =~ "([0-9, ]+) M[bB] total memory in use") :: String
        mem = case length bytes of
          0 -> 0
          _ -> (read (filter (`elem` "0123456789") bytes) :: Int)
    recordMemoryUsed $ mem * 1024 * 1024
    return res

git :: String -> [String] -> Command String
git cmd args = case args of
  ["obliterate", "--last", n, "--all"] -> obliterate n
  _ -> runInPlayground cmd ("--no-pager" : args')
  where args' | ("whatsnew":rem) <- args = "diff" : rem
              | ("revert":rem) <- args = "reset" : "--hard" : filter (/="-a") rem
              | ("record":rem) <- args = "commit" : filter (/= "--no-test") rem
              | ("get":rem) <- args = "clone" : rem
              | otherwise = args
        obliterate n = do
          rev <- (last . lines) `fmap` runInPlayground cmd
                             ["rev-list", "--max-count=" ++ n, "HEAD"]
          runInPlayground cmd ["reset", "--hard", rev]

hg :: String -> [String] -> Command String
hg cmd args = case args of
  ["obliterate", "--last", "1", "--all"] -> rollback
  _ ->  runInPlayground cmd args'
  where args' | ("whatsnew":rem) <- args = "diff" : rem
              | ("record":rem) <- args = "commit" : [ if opt == "--author" then "--user" else opt
                                                    | opt <- rem, opt /= "--no-test", opt /= "--all" ]
              | ("get":rem) <- args = "clone" : rem
              | otherwise = args
        rollback = do runInPlayground cmd ["rollback"]
                      runInPlayground cmd ["revert", "-a"]

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 -> TestBinary -> Variant -> Command ()
mkVariant origrepo bin v = do
  isrepo <- liftIO $ doesDirectoryExist (origrepo </> "_darcs")
  unless isrepo $ fail $ origrepo ++ ": Not a darcs repository!"
  case vId v of
    OptimizePristineVariant -> do
      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) `catchany` \_ -> return ()
        darcs (binCommand bin) [ "optimize", "--pristine", "--repodir", variant_repo ]
        return ()
    GitVariant -> do
      variant_exists <- liftIO $ doesDirectoryExist (variant_repo </> ".git")
      unless variant_exists $ do
        mkdir_p variant_repo
        sub $ do
          cd variant_repo
          system "git init"
          system ("darcs convert --export ../" ++ origrepo ++ " | git fast-import")
          system "git checkout"
        return ()
    HgVariant -> do
      variant_exists <- liftIO $ doesDirectoryExist (variant_repo </> ".hg")
      gitvariant <- case reverse variant_repo of
            'g':'h':'-':rest -> return $ reverse rest ++ "-git"
            _ -> fail "can't figure the git variant..."
      unless variant_exists $ do
        mkVariant origrepo bin (toVariant GitVariant)
        system $ "hg convert " ++ gitvariant ++ " " ++ variant_repo
        sub $ cd variant_repo >> system "hg checkout"
        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 bin =
  sequence_ [ mkVariant (trPath repo) bin variant
            | repo <- repos, variant <- trVariants repo ]