{-# LANGUAGE TypeOperators, ScopedTypeVariables, DeriveDataTypeable #-}
-- This library is for parse information from /proc on Linux.
module System.Linux.Proc (
-- * Types
    ProcessName,  
    ProcessState (..),
    ProcStatus (..),
    UID,
    UserName,
    UserDatabase,

-- * Methods
    procGetProcessDirs,
    procGetAllProcessStatus,
    getUserDatabase,
    ) where

import Control.Applicative
import Control.Exception
import Control.Monad
import Data.Char
import Data.Int
import Data.List
import Data.List.Split
import Data.Map (Map)
import Data.Maybe
import Data.Typeable
import Data.Word
import System.Directory
import System.Exit
import System.FilePath
import System.IO.Error
import System.Posix.Types
import System.Process
import Text.Regex.TDFA
import Text.XFormat.Read

import qualified Control.Exception as Exc
import qualified Data.Map as M
import qualified System.IO.Strict as IO

type ProcessName = String
type ProcessInfo = [String]
data ProcessState = Running
                  | Sleeping
                  | Waiting
                  | Zombie
                  | Traced
                  | Paging
                    deriving (Show, Read, Ord, Eq, Typeable)
type UID = Int
type UserName = String
type UserDatabase = Map UID UserName
type CPUTime = (Int, Int, Int, Int, Int)
type CPUTimeDatabase = Map Int (Maybe CPUTime)

data ProcStatus = 
    ProcStatus {psProcessId      :: Int          -- The process ID.
               ,psCommand        :: String       -- The filename of the executable
               ,psState          :: ProcessState -- Process state
               ,psParentProcessId:: Int          -- The PID of the parent
               ,psProcessGroupId :: Int          -- The process group ID of the process
               ,psSessionId      :: Int          -- The session ID of the process
               ,psNice           :: Int          -- The nice value from range 19 (low priority) to -20 (high priority)
               ,psNumThreads     :: Int          -- Number of threads in this process
               ,psVirtualMem     :: Int          -- peak virtual memory size in bytes
               ,psResidentMem    :: Int          -- resident size in bytes
               ,psUID            :: Int          -- UID of process
               ,psUsername       :: String       -- user name of process
               ,psCmdline        :: String       -- the complete command-line of process, unless the process is a zombie.
               ,psCpuPercent     :: Double       -- CPU percent
               } deriving (Show, Read, Ord, Eq, Typeable)

-- | Path for /proc.
procPath :: FilePath
procPath = "/proc"

-- | Status file about the process, 
procStatFile :: FilePath
procStatFile = "stat"

-- | Status file.
procStatusFile :: FilePath
procStatusFile = "status"

-- | Cmdline file.
procCmdlineFile :: FilePath
procCmdlineFile = "cmdline"

-- | Information Box to find status.
psCommandBox          = (1,   String,   "psCommandBox")
psStateBox            = (2,   Char,     "psStateBox")
psParentProcessIdBox  = (3,   Int,      "psParentProcessIdBox")
psProcessGroupIdBox   = (4,   Int,      "psProcessGroupIdBox")
psSessionIdBox        = (5,   Int,      "psSessionIdBox")
psUtimeBox            = (13,  Int,      "psUtimeBox")
psStimeBox            = (14,  Int,      "psStimeBox")
psCutimeBox           = (15,  Int,      "psCutimeBox")
psCstimeBox           = (16,  Int,      "psCstimeBox")
psNiceBox             = (18,  Int,      "psNiceBox")
psNumThreadsBox       = (19,  Int,      "psNumThreadsBox")
psProcessorBox        = (38,  Int,      "psProcessorBox")

-- | Get process directories.
procGetProcessDirs :: IO [FilePath]
procGetProcessDirs = 
    -- Get directory contents.
    getDirectoryContents procPath
    -- Filter sub-directories.
    >>= filterM (\x -> doesDirectoryExist $ procPath </> x) 
    -- Filter integer directories.
    >>= \dirs -> return $ filter isIntegerString dirs

-- | Get process id.
procGetProcessIDs :: IO [Int]
procGetProcessIDs = 
  liftM (map (\x -> read x :: Int)) procGetProcessDirs 

-- | Get process status with given process id.
procGetProcessStatus :: Int -> UserDatabase -> IO (Maybe (String, ProcessState, Int, Int, Int, Int, Int, Int, Int, Int, String, String, Int))
procGetProcessStatus pid database = do
  info <- procGetStatInfo pid
  case info of 
    Just i  -> do
       -- Get status from /proc/pid/stat file.
       let (cmd, state, ppid, pgid, sid, nice, threads, processor) = procParseStatInfo i

       -- Get command line from /proc/pid/cmdline file.
       cmdline <- procGetCmdlineInfo pid

       -- Get command name.
       let cmdStr  = head $ splitOneOf "\0 " cmdline -- split with null separated or blank character
           command = 
             -- Because name in /proc/pid/stat has number limit,
             -- so we try to use name in /proc/pid/cmdline.
             if null cmdStr
                -- Use name in /proc/pid/stat file if /proc/pid/cmdline is empty
                then cmd
                -- Otherwise take file name of cmdline.
                else takeFileName cmdStr

       -- Get status from /proc/pid/status file.
       (vSize, vRss, uid) <- procGetStatusInfo pid

       -- Find username with UID.
       let filterMap = M.filterWithKey (\ id _ -> id == uid) database
           username  = if M.null filterMap
                          then ""
                          else snd $ M.findMin filterMap

       return $ Just (command, state, ppid, pgid, sid, nice, threads, vSize, vRss, uid, username, cmdline, processor)
    Nothing -> return Nothing

-- | Get information for all running processes.
procGetAllProcessStatus :: IO [ProcStatus]
procGetAllProcessStatus = do
  -- Get process ids.
  ids <- procGetProcessIDs

  -- Get before cpu time database.
  beforeDatabase <- procGetCPUTimeDatabase ids M.empty

  -- Get user database.
  userDatabase <- getUserDatabase

  -- Get status.
  statusList <- mapM (\pid -> do
                       status <- procGetProcessStatus pid userDatabase
                       return (pid, status)
                    ) ids

  -- Get after cpu time database.
  afterDatabase <- procGetCPUTimeDatabase ids M.empty

  -- Get process status.
  let procStatusList = 
          map (\ (pid, status) -> 
                   case status of
                     Just (command, state, ppid, pgid, sid, nice, threads, vSize, vRss, uid, username, cmdline, processor) ->
                         let cpuPercent = procGetCPUPercent pid processor beforeDatabase afterDatabase
                         in Just $ ProcStatus pid command state 
                                              ppid pgid sid nice 
                                              threads vSize vRss 
                                              uid username cmdline 
                                              cpuPercent
                     Nothing -> Nothing
              ) statusList

  return $ catMaybes procStatusList

-- | Get CPU percent.
procGetCPUPercent :: Int -> Int -> CPUTimeDatabase -> CPUTimeDatabase -> Double
procGetCPUPercent pid processor beforeDatabase afterDatabase = 
    let bDatabase = (M.filterWithKey (\ id _ -> id == pid) beforeDatabase)
        aDatabase = M.filterWithKey (\ id _ -> id == pid) afterDatabase 
    in if M.null bDatabase || M.null aDatabase
          then 0
          else 
              let beforeTimes = snd $ M.findMin bDatabase
                  afterTimes  = snd $ M.findMin aDatabase
              in 
                if isNothing beforeTimes || isNothing afterTimes
                   -- Return 0 if any time database is empty.
                   then 0
                   else 
                       let (beforeUtime, beforeStime, beforeCutime, beforeCstime, beforeTotalTime) = fromJust beforeTimes
                           (afterUtime, afterStime, afterCutime, afterCstime, afterTotalTime) = fromJust afterTimes
                           -- Get process cpu time different.
                           processCPUTimeDiff = 
                               fromIntegral 
                               (sum [afterUtime, afterStime, afterCutime, afterCstime] -
                                sum [beforeUtime, beforeStime, beforeCutime, beforeCstime])
                           -- Get total cpu time different.
                           totalCPUTimeDiff = fromIntegral (afterTotalTime - beforeTotalTime)
                           -- Get cpu number.
                           cpuNumber = fromIntegral processor
                       in if totalCPUTimeDiff == 0
                             then 0
                             else formatFloatN ((processCPUTimeDiff * 100) / totalCPUTimeDiff * cpuNumber) 2

-- | Get information of /proc/pid/stat.
procGetStatInfo :: Int -> IO (Maybe ProcessInfo)
procGetStatInfo pid = do
  let filepath = procPath </> show pid </> procStatFile
  Exc.catch 
     (liftM (Just . words) $ IO.readFile filepath)
     (\(_ :: IOError) -> return Nothing) -- skip current file if failed with `readFile`

-- | Get memory information of /proc/pid/status.
procGetStatusInfo :: Int -> IO (Int, Int, Int)  
procGetStatusInfo pid = do
  let filepath = procPath </> show pid </> procStatusFile
  Exc.catch 
     (do
       lns <- liftM lines $ IO.readFile filepath
       let pickFirstInt match list = 
             let matchList = filter (\x -> match `isInfixOf` x) list
             in if null matchList
                   then 0
                   else 
                       let (_, value, _) = head matchList =~ "[0-9]+" :: (String, String, String)
                       in if null value
                             then 0
                             else read value :: Int
           size = pickFirstInt "VmSize" lns
           rss  = pickFirstInt "VmRSS" lns
           uid  = pickFirstInt "Uid" lns
       return (1024 * size, 1024 * rss, uid) -- convert Kb to byte
     )
     (\(_ :: IOError) -> do
        putStrLn $ "procGetStatusInfo : read " ++ filepath ++ " failed."
        return (0, 0, 0))

-- | Get information from /proc/pid/cmdline.
procGetCmdlineInfo :: Int -> IO String
procGetCmdlineInfo pid = do
  let filepath = procPath </> show pid </> procCmdlineFile
  Exc.catch
     (IO.readFile filepath)
     (\(_ :: IOError) -> do
          putStrLn $ "procGetCmdlineInfo : read " ++ filepath ++ " failed."
          return "")

-- | Parse status information from /proc/pid/stat.
procParseStatInfo :: ProcessInfo -> (String, ProcessState, Int, Int, Int, Int, Int, Int)
procParseStatInfo info = 
    (toProcessName $ procPickStat info psCommandBox 
    ,toProcessState $ procPickStat info psStateBox 
    ,procPickStat info psParentProcessIdBox
    ,procPickStat info psProcessGroupIdBox
    ,procPickStat info psSessionIdBox
    ,procPickStat info psNiceBox
    ,procPickStat info psNumThreadsBox
    ,procPickStat info psProcessorBox)

-- | Get database of cpu time.
procGetCPUTimeDatabase :: [Int] -> CPUTimeDatabase -> IO CPUTimeDatabase
procGetCPUTimeDatabase [] database = return database
procGetCPUTimeDatabase (x:xs) database = do
  cpuTime <- procGetProcessCPUTime x
  procGetCPUTimeDatabase xs $ M.insert x cpuTime database

-- | Parse status time information from /proc/pid/stat.
procGetProcessCPUTime :: Int -> IO (Maybe (Int, Int, Int, Int, Int))
procGetProcessCPUTime pid = do
  info <- procGetStatInfo pid
  case info of 
    Just i -> do
        cpuTotalTime <- getTotalCPUTime
        return $ Just  
          (procPickStat i psUtimeBox
          ,procPickStat i psStimeBox
          ,procPickStat i psCutimeBox
          ,procPickStat i psCstimeBox
          ,cpuTotalTime)
    Nothing -> 
      return Nothing

-- | Find status from given info.
-- Throw error if parse failed.
procPickStat :: Format d a => ProcessInfo -> (Int, d, String) -> a 
procPickStat infos (index, format, debugInfo) = 
  maybeError (case infos ?! index of
                Just status -> readf format status
                Nothing     -> Nothing)
             $ "parse error : " ++ show debugInfo

-- | To Process name.
toProcessName :: String -> String
toProcessName = filter (`notElem` "()")

-- | To process state.
toProcessState :: Char -> ProcessState
toProcessState 'R'   = Running
toProcessState 'S'   = Sleeping
toProcessState 'D'   = Waiting
toProcessState 'Z'   = Zombie
toProcessState 'T'   = Traced
toProcessState 'W'   = Paging
toProcessState state = error $ "Miss state : " ++ show state

-- | Indicate error or return a.
-- This function for replace `fromJust`, expression `fromJust x` is bad
-- when `x` is `Nothing`, so `maybeError` allowed you customize error
-- informatiron.
maybeError :: Maybe a -> String -> a
maybeError m str = fromMaybe (error str) m

-- | Return element of list with given index.
(?!) :: [a] -> Int -> Maybe a
[] ?! _ = Nothing              
xs ?! n
    | n < 0 || n >= length xs 
        = Nothing
    | otherwise               
        = listToMaybe . drop n $ xs

-- | Is integer string? 
isIntegerString :: String -> Bool
isIntegerString []     = False
isIntegerString [x]    = isDigit x
isIntegerString (x:xs) = isDigit x && isIntegerString xs

-- | Get user database for search username with UID.
getUserDatabase :: IO UserDatabase
getUserDatabase = do
  (exitCode, result, _) <- readProcessWithExitCode "getent" ["passwd"] ""
  return $ case exitCode of
             ExitFailure _ -> M.empty
             ExitSuccess   -> 
               M.fromList $ 
                map (\str -> 
                         let (username, rest) = break (== ':') str
                             (_, uid, _) = rest =~ "[0-9]+" :: (String, String, String)
                         in (read uid :: Int, username)
                    ) $ lines result

-- | Get total cpu time.
getTotalCPUTime :: IO Int
getTotalCPUTime = do
  let filepath = procPath </> procStatFile
  Exc.catch 
     (do
       list <- liftM lines $ IO.readFile filepath
       if null list
          then return 0
          else do
            let strList = tail $ words $ head list
                num = sum $ map (\x -> read x :: Int) strList
            return num)
     (\(_ :: IOError) -> return 0)

-- | Format float with specify precision.
formatFloatN :: Double -> Int -> Double
formatFloatN n d =
  fromIntegral (floor $ n * 10 ^ d) / 10 ^ d