--
-- Copyright (c) 2006 Don Stewart - http://www.cse.unsw.edu.au/~dons
-- GPL version 2 or later (see http://www.gnu.org/copyleft/gpl.html)
--

--
-- darcs-graph:
--  generate .png graphs of commit activity of a darcs repository,
--  including a 7 day sliding average
--
-- requires gnuplot
--

import Data.List
import Data.Maybe
import qualified Data.IntMap as I

import Data.ByteString.Char8 (ByteString,pack,unpack)
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString as B

import Control.Exception        (bracket,handle)
import Control.Monad            (when)

import System.Console.GetOpt
import System.Environment       (getArgs)
import System.IO
import System.IO.Unsafe
import System.Directory
import System.Process
import System.Exit
import Data.Time
import System.Locale

tmp :: FilePath
tmp = unsafePerformIO $ getTemporaryDirectory

data Flag = DateFlag String | YFlag Int | Help | OFile FilePath | PName String
    deriving (Show,Eq)

options :: [OptDescr Flag]
options = [
  Option ['f'] ["filter"] (ReqArg (\s -> DateFlag s) "DATE")
        "Don't display patches before date. Date is in format \"yyyymmdd\"",
  Option ['y'] ["y"]      (ReqArg (\s -> YFlag (read s)) "INT")
        "Vertical scale of graphs. Integer.",
  Option ['o'] ["output"] (ReqArg (\s -> OFile s) "FILE")
        "Filename to write output to.",
  Option ['n'] ["name"]   (ReqArg (\s -> PName s) "NAME")
        "Name of project used in generated graph.",
  Option ['h'] ["help"]   (NoArg Help)
        "Display this help"
  ]

usage :: String
usage = usageInfo "usage: darcs-graph [options] repo\n" options

readOptions :: (Monad m) => [String] -> m ([Flag], [String])
readOptions args =
  case getOpt Permute options args of
    (o, n, [])    -> return (o, n)
    (_, _, errs)  -> fail $ unlines errs

getFilter :: [Flag] -> Maybe Day--ByteString
getFilter l =
  let l' = [s | (DateFlag s) <- l]
  in if null l' then Nothing else parseTime defaultTimeLocale "%Y%m%d" (head l')

getYFlag :: [Flag] -> Maybe Int
getYFlag l =
  let l' = [i | (YFlag i) <- l]
  in if null l' then Nothing else Just . head $ l'

getOFile :: [Flag] -> String -> FilePath
getOFile flags proj = case [ s | OFile s <- flags ] of
    [] -> tmp </> proj ++ "-commits.png"
    (fn:_) -> fn

getPName :: [Flag] -> String -> String
getPName flags proj = case [ s | PName s <- flags ] of
    [] -> proj
    (fn:_) -> fn

data Range = Small | Big | Normal
    deriving Eq

-- let's go
main :: IO ()
main = do
    args <- getArgs
    (flags, remaining) <- readOptions args
    when (null remaining || length remaining > 1 || Help `elem` flags) $ error usage
    let path    = pack . head $ remaining
        mfilter = getFilter flags
        mymax   = getYFlag flags

    -- aquire some temp files, make sure to clean them up afterwards
    bracket
        (do x <- openTempFile tmp "darcs-graph.data"
            y <- openTempFile tmp "darcs-graph-average.data"
            return (x,y))
        (\((f,_),(g,_)) -> removeFile f >> removeFile g) $ \((f,h),(g,i)) -> do

    changes <- darcsChanges (unpack path)

    let stats' = uniq $ sort $ parseChanges changes

    -- filter out stats earlier than 'yyyymmdd'
    let stats  | (Just date) <- mfilter = filter (\(x,_) -> x > date) stats'
               | otherwise              = stats'

    when (null stats) $ error $ "no patches found in: "++unpack path
    (C.hPutStr h . C.unlines . map fmt) stats >> hClose h

    -- now work out the sliding average, and write it out too
    let days = enumDates (fst . head $ stats) (fst . last $ stats)
        avrg = sliding (map readFst stats) (map toInt days)
    (C.hPutStr i . C.unlines . map fmt') avrg >> hClose i

    -- now run gnuplot on the result
    let proj   = (unpack . basename $ path)
        out    = getOFile flags proj
        script = gnuplotScript f    -- handle small repos
                    -- can't generate average with less than 7 points,
                    -- and the sliding line needs 3 points to plot csplines
                    (if length days >= window && length avrg >= 3
                        then Just g
                        else Nothing)

                    (if length days < 120
                        then Small
                        else if length days > 2500
                            then Big
                            else Normal) -- only print days or months if range is ok
                    (getPName flags proj)
                    mymax
                    out


    (ih,_,eh,pid) <- handle
        (\e -> error $ "Couldn't fork gnuplot: " ++ show e)
        (do mgnu <- findExecutable "gnuplot"
            case mgnu of
                Nothing      -> error "Cannot find gnuplot"
                Just gnuplot -> runInteractiveCommand gnuplot)

    hPutStrLn ih script >> hClose ih -- write into gnuplot
    hGetContents eh >>= putStr       -- dump any error output it produces
    waitForProcess pid

    putStrLn $ "Output written to: " ++ out

    where
      readFst (a,b) = (toInt a, b)
      uniq          = map (\n -> (head n, length $ n)) . group

      fmtCT c    = formatTime defaultTimeLocale "%Y%m%d" c
      fmt (s,n)  = joinWithSpace (fmtCT s)               (show n)
      fmt' (n,d) = joinWithSpace (show n) (show d)

      joinWithSpace s t = B.intercalate (B.singleton 32) [pack s,pack t]

parseChanges :: ByteString -> [Day]
parseChanges chs
  = worker (C.lines chs)
  where worker [] = []
        worker (c:cs) | Just utcTime <- parseTime defaultTimeLocale "%a %b %e %X %Z %Y" (unwords $ take 6 $ words $ C.unpack c)
          = utcTime:worker cs
        worker (_:cs) = worker cs

darcsChanges :: FilePath -> IO ByteString
darcsChanges path
  = do (inh,outh,errh,pid) <- runInteractiveProcess "darcs" ["changes","--repo="++path] Nothing Nothing
       hClose inh
       hClose errh
       out <- B.hGetContents outh
       code <- waitForProcess pid
       case code of
         ExitFailure c -> do hPutStrLn stderr "`darcs --changes` failed."
                             exitWith (ExitFailure c)
         ExitSuccess   -> return out

------------------------------------------------------------------------
-- Dealing with clock times

--
-- enumerate all the days, as ClockTimes, between the two dates
--
enumDates :: Day -> Day -> [Day]
enumDates d0 dn = map (\d -> addDays d d0) [0 .. days]
    where
        days    = diffDays dn d0

--
-- convert a ClockTime back to an Int in yyyymmdd format
--
toInt :: Day -> Int
toInt ct = read $ formatTime defaultTimeLocale "%Y%m%d" ct

-- sliding average window
window :: Int
window = 7

--
-- compute the 7-day sliding average, starting with a sparse sequence
--
sliding :: [(Int,Int)] -> [Int] -> [(Int,Double)]
sliding [] _    = []
sliding xs days = [ (i,slidingAv i) | i <- days' ]
  where
    m     = window `div` 2 -- 1/2 n

    days' = drop m . take (max 0 (length days - window+m+1)) $ days
    table = I.fromList xs :: I.IntMap Int -- to handle missing elems easily

    slidingAv :: Int -> Double
    slidingAv i = fromIntegral (sum seqN) / fromIntegral window
        where seqN = map (\j -> I.findWithDefault 0 j table) [i-m..i+(window-m-1)]

--
-- generate the gnuplot script. only generate the avergage line, if
-- there repo day range is larger than a window
--
gnuplotScript :: FilePath -> Maybe FilePath -> Range -> String -> Maybe Int -> FilePath -> String
gnuplotScript f mg range p mymax outFile = header ++ case mg of
    Nothing -> "plot \""++f++"\" using 1:2 with points pointsize 0.5 pointtype 2\n"
    Just g  -> "plot \"" ++
                    f++"\" using 1:2 with points pointsize 0.5 pointtype 2, \"" ++
                    g++"\" using 1:2 smooth csplines 2\n"

  where header = unlines
           ["set xdata time"
           ,"set timefmt \"%Y%m%d\""
           ,case range of
                Small  -> "set format x\"%y%m%d\""
                Normal -> "set format x\"%y%m\""
                Big    -> "set format x\"%y\""
           ,"set size 1,0.4"
           ,"set mxtics 0"
           ,"set xtics nomirror"
           ,"set ytics nomirror"
           ,"set border 3"
           ,"set yrange [0:" ++ (if isNothing mymax then "*" else show . fromJust $ mymax) ++ "]"
           ,"set title \""++p++" commits per day and 7-day sliding average\""
           ,"unset key"
           ,"set ylabel \"commits per day\""

-- To use a TTF font, set the GDFONTPATH=/home/dons/TTF, and uncomment
-- this line:
--         ,"set terminal png font Vera 9 xffffff x000000 x404040 xd0d0d0 x0000ff"

           ,"set terminal png small xffffff x000000 x404040 xd0d0d0 x0000ff"
           ,"set output \"" ++ outFile ++ "\""]

--
-- ByteString utilities
--

-- | Packed string version of basename
basename :: ByteString -> ByteString
basename fps = case C.elemIndexEnd '/' s of
    Nothing -> fps
    Just i  -> C.drop (i+1) s
    where s = C.reverse . C.dropWhile (=='/') . C.reverse $ fps

-- useful
infixr 6 </>

(</>) :: FilePath -> FilePath -> FilePath
[] </> b = b
a  </> b = a ++ "/" ++ b