#!/usr/bin/runhaskell

ps2xdu -- short program for visualizing results of ps command for identifying which programs (processes and their descendants) use lots of CPU time and memory resource.
This program converts results of ps command into the format that the xdu command understands, and invokes xdu.
(So requires xdu.)

\begin{code}
import System.Console.GetOpt
import System.Process(readProcess, runProcess, waitForProcess)
import System.IO
import System.Directory(removeFile)
import System.Environment
import System.Exit
import qualified Data.IntMap as IM
import Data.Maybe(maybeToList)
import Data.Function(on)
import Data.List(sortBy)

data Flag = C | VSZ | RSS | Keys String | Output String -- | Qualify | Unqualify
options :: [OptDescr Flag]
options = [ Option ['c'] ["c","cpu"]     (NoArg C) "extract the CPU utilization"
          , Option ['v'] ["vsz","vsize"] (NoArg VSZ) "extract the virtual memory size (default)"
          , Option ['r'] ["rss","rsz","rssize"] (NoArg VSZ) "extract the resident set size"
          , Option ['k'] ["keys"]        (ReqArg Keys "KEYS") "comma-separated keys to be shown, which will be passed to `ps -o'. (default=`cmd,uid')"
          , Option ['o'] []              (ReqArg Output "FILE") "output to FILE. If FILE=='-' stdout is chosen."
--          , Option ['q'] ["qualify"]          (NoArg Qualify) "qualify variable names shown." 
--          , Option ['u'] ["unqualify"]       (NoArg Unqualify) "unqualify variable names shown. This is useful if you are using a small display, though it does not work correctly when there are remarkable cost centers with the same name in different modules." 
          ]
readOpts :: IO ([Flag], [String])
readOpts = do
	      argv     <- getArgs
	      case (getOpt Permute options argv) of
			    (o,n,[]  ) -> return (o,n)
			    (_,_,errs) -> do hPutStrLn stderr (concat errs)
                                             usage
                                             exitFailure
usage :: IO ()
usage = do progname <- getProgName
           hPutStrLn stderr $ usageInfo ("Usage: "++progname++" [OPTION...]") options

data Stat = St {ext :: String, keys :: String, toFile :: Maybe String} -- , qual :: Bool}
defaultStat = St {ext="vsz", keys="cmd", toFile=Nothing} -- , qual=True}
procFlags :: [Flag] -> Stat
procFlags = foldl procFlag defaultStat
procFlag :: Stat -> Flag -> Stat
procFlag st C            = st{ext="c"}
procFlag st VSZ          = st{ext="vsz"}
procFlag st RSS          = st{ext="rss"}
procFlag st (Keys str)   = st{keys=str}
procFlag st (Output str) = st{toFile=Just str}
-- procFlag st Qualify      = st{qual=True}
-- procFlag st Unqualify    = st{qual=False}

unwordsBy c = tail . unwordsBy' c
unwordsBy' c = concat . map (c:)

type PID = Int

data Tree = T {rootPID::PID, total::Int, rootKeys::String, children::[Tree]} deriving Show

mkForest :: PID -> IM.IntMap [(PID, Int, String)] -> [Tree]
mkForest ppid im = [ T {rootPID = pid, total = v + sum (map total forest), rootKeys = str, children = forest} | ts <- maybeToList $ IM.lookup ppid im, (pid, v, str) <- ts, let forest = mkForest pid im ]

treeToReverseDU :: Tree -> [(Int, String)]
treeToReverseDU (T _ 0 _ _) = []
treeToReverseDU (T p v s c)  = (v,s) : sortBy (compare `on` fst) [(cv, s++'/':cs) | t <- c, (cv, cs) <- treeToReverseDU t]

removeCmdPath = (\a b -> reverse (takeWhile (/= a) (reverse b)))

main = do (flags,args) <- readOpts
          let stat = procFlags flags
          psResult <- readProcess "ps" ["-e", "-o", "ppid,pid,"++ext stat++',':keys stat] []
          let fromPPID = IM.fromListWith (++) $ map ((\ (ppid:pid:v:keys) -> (read ppid, [(read pid, read v, pid++unwordsBy' '_' (map (removeCmdPath '/') keys))])). words) $ tail $ lines psResult
          let result = unlines $ reverse [ shows v $ ' ':keys | tr <- mkForest 0 fromPPID, (v, keys) <- treeToReverseDU tr ]
          -- The rest is the same as prof2xdu.
          case toFile stat of Nothing -> do (filepath,handle) <- openTempFile "." "p2x.tmp"
                                            hPutStr handle result
                                            hClose handle
                                            handle <- openFile filepath ReadMode
                                            ph <- runProcess "xdu" [] Nothing Nothing (Just handle) Nothing Nothing
                                            waitForProcess ph
                                            removeFile filepath
                              Just "-" -> putStr result
                              Just fn  -> writeFile fn result
\end{code} 
 readProcess "pstree" ["-lAp"]  --  This was not used.