--
-- 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,append,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 Foreign.C.Error
import Foreign.Storable         (peek)
import Foreign.Ptr              (Ptr, nullPtr)
import Foreign.Marshal          (alloca)

import System.Console.GetOpt
import System.Environment       (getArgs)
import System.Posix.Internals
import System.IO.Error          (modifyIOError, ioeSetFileName)
import System.IO
import System.IO.Unsafe
import System.Time
import System.Directory
import System.Process

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 ByteString
getFilter l =
  let l' = [s | (DateFlag s) <- l]
  in if null l' then Nothing else Just . pack . 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

    -- read patches dir, count patch dates, and write to temp file
    ps <- readDir (path `append` pack "/_darcs/patches/")

    let stats' = uniq . sort . map (C.take 8) . filter ignore $ ps

    -- 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 (toClock . fst . head $ stats) (toClock . 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) = ((fst . fromJust . C.readInt) a, b)
      uniq          = map (\n -> (head n, length $ n)) . group

      fmt (s,n)  = joinWithSpace s               (pack . show $ n)
      fmt' (n,d) = joinWithSpace (pack.show $ n) (pack.show $ d)

      ignore f = f /= pack "pending"
              && f /= pack "unrevert"
              && not (pack "." `C.isPrefixOf` f)

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


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

--
-- enumerate all the days, as ClockTimes, between the two dates
--
enumDates :: ClockTime -> ClockTime -> [ClockTime]
enumDates d0 dn = map (\d -> addToClockTime (TimeDiff 0 0 d 0 0 0 0) d0) [0 .. days]
    where
        t       = diffClockTimes dn d0
        days    = floor (((fromIntegral $ tdSec t) / 60 / 60 / 24) :: Double)

--
-- convert a date packed in a bytestring to a ClockTime
-- yyyymmdd format.
--
toClock :: ByteString -> ClockTime
toClock x = toClockTime $ CalendarTime {
                ctYear   = y,
                ctMonth  = toEnum (m-1) :: Month,
                ctDay    = d,
                ctHour   = 0, ctMin = 0, ctSec = 0, ctPicosec = 0,
                ctWDay   = undefined, ctYDay = undefined,
                ctTZName = "GMT", ctTZ = 0 , ctIsDST = False
            }
    where (y,m,d) = split x

          -- read the bytestring as a triple of yyyymmdd
          split s = let a = fst . fromJust . C.readInt . C.take 4 $ s
                        b = fst . fromJust . C.readInt . C.drop 4 . C.take 6 $ s
                        c = fst . fromJust . C.readInt . C.drop 6 . C.take 8 $ s
                    in (a,b,c)

--
-- convert a ClockTime back to an Int in yyyymmdd format
--
toInt :: ClockTime -> Int
toInt ct = read $ show yyyy ++ show0 mm ++ show0 dd
    where yyyy = ctYear cl
          mm   = fromEnum (ctMonth cl) + 1
          dd   = ctDay cl

          cl   = unsafePerformIO (toCalendarTime ct)
          show0 n | n < 10  && n >= 0 = '0' : show n
                  | otherwise         = show n

-- 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

--
-- | Packed version of get directory contents. super fast
--
readDir :: ByteString -> IO [ByteString]
readDir path = do
  modifyIOError (`ioeSetFileName` (unpack path)) $
   alloca $ \ ptr_dEnt ->
     bracket
    (C.useAsCString path $ \s ->
       throwErrnoIfNullRetry desc (c_opendir s))
    (\p -> throwErrnoIfMinus1_ desc (c_closedir p))
    (\p -> loop ptr_dEnt p)
  where
    desc = "readDir"

    loop :: Ptr (Ptr CDirent) -> Ptr CDir -> IO [ByteString]
    loop ptr_dEnt dir = do
      resetErrno
      r <- readdir dir ptr_dEnt
      if (r == 0)
        then do dEnt <- peek ptr_dEnt
                if (dEnt == nullPtr)
                    then return []
                    else do  -- copy entry out before we free:
                        entry <- C.packCString =<< d_name dEnt
                        C.length entry `seq` return ()  -- strictify
                        freeDirEnt dEnt
                        entries <- loop ptr_dEnt dir
                        return $! (entry:entries)

        else do errno <- getErrno
                if (errno == eINTR)
                    then loop ptr_dEnt dir
                    else do let (Errno eo) = errno
                            if (eo == end_of_dir)
                                then return []
                                else throwErrno desc

-- useful
infixr 6 </>

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