module Benchmark where

import Shellish hiding ( run )
import Control.Applicative
import Data.Char
import Data.List
import Data.Maybe
import System.Directory
import System.FilePath( (</>) )
import System.IO
import qualified Text.Tabular          as Tab
import TabularRST as TR
import System.Exit
import Text.Printf
import Text.Regex.Posix( (=~) )
import Data.Time.Clock
import Control.Monad.Error
import Control.Monad.State( liftIO )
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 Text.JSON

precision, iterations :: Int
precision = 1
iterations = 2
combine :: Ord a => [a] -> a
combine = minimum

type Darcs = [String] -> Command String
type BenchmarkCmd a = Darcs -> TestRepo -> Command a

data MemTime = MemTime Rational Float deriving (Read, Show, Ord, Eq)
data TestRepo = TestRepo { trName     :: String
                         , trPath     :: FilePath -- ^ relative to the config file
                         , trAnnotate :: Maybe FilePath -- ^ relative to repo, eg. @Just "README"@
                         }
 deriving (Read, Show, Eq, Ord)

instance JSON TestRepo where
  readJSON (JSObject o) =
     TestRepo <$> jlookup "name" <*> jlookup "path"
              <*> jlookupMaybe "annotate"
    where
     jlookup a = case lookup a (fromJSObject o) of
                   Nothing -> fail "Unable to read TestRepo"
                   Just v  -> readJSON v
     jlookupMaybe a = case lookup a (fromJSObject o) of
                        Nothing     -> return Nothing
                        Just JSNull -> return Nothing
                        Just v -> Just <$> readJSON v
  readJSON _ = fail "Unable to read TestRepo"
  showJSON = error "showJSON not defined for TestRepo yet"

data TestBinary = TestBinary String deriving (Show, Read)

data Benchmark a = Idempotent String (BenchmarkCmd a)
                 | Destructive String (BenchmarkCmd a)
                 | Description String

instance Show (Benchmark a) where
    show (Idempotent s _) = s
    show (Destructive s _) = s
    show (Description s) = s

instance Read (Benchmark a) where
    readsPrec _ str = [(Description str, "")]

data Test a = Test (Benchmark a) TestRepo TestBinary deriving (Read, Show)

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)

description :: Benchmark a -> String
description (Idempotent d _) = d
description (Destructive d _) = d
description (Description d) = d

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 :: Test a -> Command (Maybe MemTime)
run (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 <- sequence [
                    do progress (show i) >> verbose ("# try " ++ show i)
                       sub $ do prepareIfDifferent (trPath tr)
                                timed (exec benchmark darcs_path tr)
                           | i <- [1 .. iterations] ]
          let time = combine [ t | MemTime _ t <- times ]
              mem = combine [ m | MemTime m _ <- times ]
              spaces = 45 - (length bin + length (description benchmark) + length (trName tr))
              result = MemTime mem time
              result_str = (concat $ intersperse ", " $ formatResult result)
          progress $ (replicate spaces ' ') ++ result_str ++ "\n"
          verbose $ "# result: " ++ result_str
          return result

formatNumber :: (PrintfArg a, Fractional a) => a -> String
formatNumber = printf $ "%."++(show precision)++"f"

formatResult :: MemTime -> [String]
formatResult (MemTime mem time) =
  [ formatNumber time ++ "s, " ++ formatNumber ((realToFrac (mem / (1024*1024))) :: Float) ++ "M" ]

tabulateRepo :: String -> [(Test a, Maybe MemTime)] -> Tab.Table String String String
tabulateRepo repo results = Tab.Table rowhdrs colhdrs rows
 where
  rowhdrs = Tab.Group Tab.NoLine $ map Tab.Header rownames
  colhdrs = Tab.Group Tab.SingleLine $ map Tab.Header colnames
  colnames = nub [ label | (Test _ _ (TestBinary label), _) <- interesting ]
  rownames = nub [ description bench | (Test bench _ _, _) <- interesting ]
  interesting = [ test | test@(Test _ r _, _) <- results, trName r == repo ]
  rows = [ concat [ fmt $ find (match row column) interesting | column <- colnames ]
           | row <- rownames ]
  match bench binary (Test bench' _ (TestBinary binary'), _) =
      bench == description bench' && binary == binary'
  fmt (Just (_, Just x)) = formatResult x
  fmt _ = [ "-, -" ]

tabulate :: [(Test a, Maybe MemTime)] -> [(String, Tab.Table String String String)]
tabulate results = zip repos $ map (flip tabulateRepo results) repos
 where repos = nub [ trName r | (Test _ r _, _) <- results ]

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)

check_darcs :: String -> IO ()
check_darcs cmd = do
       (_,outH,_,procH) <- runInteractiveCommand $ cmd ++ " --version"
       out <- strictGetContents outH
       waitForProcess procH
       case out of
         '2':'.':_ -> return ()
         _ -> fail $ cmd ++ ": Not darcs 2.x binary."

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

progress :: String -> Command ()
progress m = liftIO $ do loud <- isLoud
                         unless loud $ 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

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
       (_,outH,errH,procH) <- runInteractiveProcess cmd args Nothing Nothing
       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

benchMany :: [TestRepo] -> [TestBinary] -> [Benchmark a] -> Command [(Test a, Maybe MemTime)]
benchMany repos bins benches =
    sequence [ do let test = Test bench repo bin
                  memtime <- run test
                  return (test, memtime)
               | repo <- repos, bin <- bins, bench <- benches ]

renderMany :: [(Test a, Maybe MemTime)] -> Command ()
renderMany t = sequence_ [ do echo r
                              echo $ replicate (length r) '-' ++ "\n"
                              echo_n $ TR.render id id id tab
                           | (r, tab) <- tabulate t ]