module Benchmark where

import Shellish hiding ( run )
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 qualified Text.Tabular.AsciiArt as TA
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 System.Process( runInteractiveProcess, runInteractiveCommand,
                       waitForProcess )

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

data MemTime = MemTime Rational Float
type Darcs = [String] -> Command String

newtype TestRepo = TestRepo String deriving Eq
data TestBinary = TestBinary String

type BenchmarkCmd a = Darcs -> Command a
data Benchmark a = Idempotent String (BenchmarkCmd a)
                 | Destructive String (BenchmarkCmd a)

data Test a = Test (Benchmark a) TestRepo TestBinary

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

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

exec :: Benchmark a -> FilePath -> Command a
exec (Idempotent _ cmd) darcs_path = do
  cd "_playground"
  cmd (darcs darcs_path)
exec (Destructive _ cmd) darcs_path = do
  cd "_playground"
  let cleanup = cd ".." >> rm_rf "_playground"
  res <- cmd (darcs darcs_path) `catchError` \e -> (cleanup >> throw e)
  cleanup
  return res

defaultrepo :: FilePath -> FilePath
defaultrepo path = (path </> "_darcs" </> "prefs" </> "defaultrepo")

prepare :: String -> Command ()
prepare repo = do
  echo_n "!"
  rm_rf "_playground"
  echo_n "."
  liftIO $ createDirectory "_playground"
  let playrepo = "_playground" </> "repo"
      origrepo = "repo" <.> repo
  isrepo <- liftIO $ doesDirectoryExist (origrepo </> "_darcs")
  unless isrepo $ fail $ origrepo ++ ": Not a darcs repository!"
  liftIO $ copyTree origrepo playrepo
  echo_n "."
  wd <- pwd
  liftIO $ writeFile (defaultrepo playrepo) (wd </> origrepo)

prepareIfDifferent :: String -> Command ()
prepareIfDifferent repo = do
  let playrepo = "_playground" </> "repo"
      origrepo = "repo" <.> 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 echo_n "..."
                                           else prepare repo

run :: Test a -> Command (Maybe MemTime)
run (Test benchmark (TestRepo testrepo) (TestBinary bin)) = do
  (Just `fmap` run') `catchError` \e ->
      do echo $ " error: " ++ show e
         return Nothing
  where run' = do
          echo_n $ bin ++ " " ++ description benchmark ++ " [" ++ testrepo ++ "]: "
          exe <- which $ bin
          darcs_path <- case exe of
                          Nothing -> canonize bin
                          Just p -> return p
          times <- sequence [
                    do echo_n $ show i
                       sub $ do prepareIfDifferent testrepo
                                timed (exec benchmark darcs_path)
                           | i <- [1 .. iterations] ]
          let time = combine [ t | MemTime _ t <- times ]
              mem = combine [ m | MemTime m _ <- times ]
              spaces = 45 - (length bin + length (description benchmark) + length testrepo)
              result = MemTime mem time
          echo $ (replicate spaces ' ') ++ (concat $ intersperse ", " $ formatResult result)
          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 colgrp colnames
  colgrp x = Tab.Group Tab.NoLine [Tab.Header x, Tab.Header ""]
  colnames = nub [ label | (Test _ _ (TestBinary label), _) <- interesting ]
  rownames = nub [ description bench | (Test bench _ _, _) <- interesting ]
  interesting = [ test | test@(Test _ (TestRepo r) _, _) <- results, 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 [ repo | (Test _ (TestRepo repo) _, _) <- 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."

darcs :: String -> [String] -> Command String
darcs cmd args' = do
    (res, _, stats) <- liftIO $ do
       let args = args' ++ ["+RTS", "-sdarcs-stats", "-RTS"]
       (_,outH,errH,procH) <- runInteractiveProcess cmd args Nothing Nothing
       res <- strictGetContents outH
       errs <- strictGetContents errH
       ex <- waitForProcess procH
       stats <- readFile "darcs-stats" `catch` \_ -> return ""
       case ex of
         ExitSuccess -> return ()
         ExitFailure n -> fail $ "darcs failed with error code "
                            ++ show n ++ "\nsaying: " ++ errs
       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 $ "\n=== " ++ r ++ " ===\n"
                              echo_n $ TA.render id id id tab
                           | (r, tab) <- tabulate t ]