  Copyright 2014 Mihaly Barasz & Gergely Risko
--
  Licensed under the Apache License , Version 2.0 ( the " License " ) ;
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     -2.0
--
-- Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an " AS IS " BASIS ,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.


# LANGUAGE TemplateHaskell #
{-# LANGUAGE OverloadedStrings #-}
# LANGUAGE FlexibleContexts #

module Main (main) where

import Control.Applicative
import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.STM
import Control.Lens
import Control.Monad
import qualified Data.Array.MArray as MA
import Data.Attoparsec.ByteString.Char8 (decimal, (.*>), skipSpace, skipWhile)
import qualified Data.Attoparsec.ByteString as AP
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Char
import Data.Sequence (Seq)
import qualified Data.Sequence as S
import Data.Word (Word32)
import qualified Graphics.UI.Gtk as Gtk
import HFlags
import Pipes (Producer, yield, (>->), await, Consumer, Pipe, lift, runEffect)
import System.Environment
import System.IO
import System.Posix.Signals
import Text.Printf


defineFlag "show_cpu" True "Show the CPU load indicator/icon."
defineFlag "show_mem" True "Show the memory usage indicator/icon."
defineFlag "update_interval" (500 :: Int)
  "Update interval for all indicators in milliseconds."

$(return [])

--------------------------------------------------------------------------------
  Acquiring raw CPU and RAM numbers

data CPUCounters =
  CPUCounters { cpuUser :: {-# UNPACK #-} !Int
              , cpuAllSys :: {-# UNPACK #-} !Int
              , cpuIdle :: {-# UNPACK #-} !Int
              } deriving (Show)

-- The counters are in milliseconds and cumulative since system start.
readCPUCounters :: IO CPUCounters
readCPUCounters = parseProcStat <$> readProcFile "/proc/stat"
  where
    parseProcStat = forceParse parseFirstLine
    parseFirstLine = wrapResult <$ "cpu " .*> skipSpace
                     <*> decimal <* skipSpace <*> decimal <* skipSpace
                     <*> decimal <* skipSpace <*> decimal <* skipSpace
                     <*> decimal <* skipSpace <*> decimal <* skipSpace
                     <*> decimal
    wrapResult user nice system idle iowait irq softirq =
      CPUCounters (user+nice) (system+iowait+irq+softirq) idle

data MEMCounters =
  MEMCounters { ramTotal :: {-# UNPACK #-} !Int
              , ramCached :: {-# UNPACK #-} !Int
              , ramFree :: {-# UNPACK #-} !Int
              } deriving (Show)

  The counters are in kB.
readMEMCounters :: IO MEMCounters
readMEMCounters = parseProcMeminfo <$> readProcFile "/proc/meminfo"
  where
    parseProcMeminfo = forceParse parseHead
    parseHead :: AP.Parser MEMCounters
    parseHead =
      wrapResult <$ "MemTotal:" .*> skipSpace <*> decimal <* kb <*
        "MemFree:" .*> skipSpace <*> decimal <* kb <*>
        maybeMemAvailable <*
        "Buffers:" .*> skipSpace <* (skipWhile isDigit) <* kb <*
        "Cached:" .*> skipSpace <*> decimal <* kb
    kb = skipSpace *> "kB" .*> skipSpace
    maybeMemAvailable =
      Just <$ "MemAvailable:" .*> skipSpace <*> decimal <* kb <|> pure Nothing
    wrapResult total free mAvailable cached =
      MEMCounters total cached $ maybe free id mAvailable

forceParse :: AP.Parser a -> BS.ByteString -> a
forceParse parser = either error id . AP.parseOnly parser

-- If you stat /proc/stat or /proc/meminfo they look 0 sized, so we
-- cannot use BS.readFile.
  We assume that all files of interest fit in 4k .
readProcFile :: FilePath -> IO ByteString
readProcFile file = withBinaryFile file ReadMode $ flip BS.hGetSome 4096

--------------------------------------------------------------------------------
-- Pipes

  TODO(klao ): use strict pairs
type Measurement = (Double, Double)

measureMEM :: Producer Measurement IO r
measureMEM = forever $ do
  MEMCounters total cached free <- lift readMEMCounters
  yield (fromIntegral (total-cached-free) / fromIntegral total,
         fromIntegral cached / fromIntegral total)

measureCPU :: Producer Measurement IO r
measureCPU = lift readCPUCounters >>= go
  where
    go (CPUCounters user0 sys0 idle0) = do
      m@(CPUCounters user sys idle) <- lift readCPUCounters
  TODO(klao ): handle overflows !
      let total = fromIntegral $ user + sys + idle - user0 - sys0 - idle0
      when (total > 0) $ yield (fromIntegral (user - user0) / total,
                                fromIntegral (sys - sys0) / total)
      go m

delay :: Pipe a a IO r
delay = execU (threadDelay $ flags_update_interval * 1000)

prettyPrint :: String -> Consumer Measurement IO r
prettyPrint label = forever $ do
  (x1, x2) <- await
  lift $ printf "%s: %4.1f%% %4.1f%%\n" label (x1*100) (x2*100)

collectN :: Int -> Pipe Measurement (Seq Measurement) IO r
collectN n = go S.empty
  where
    go l = do
      m <- await
      let l' = S.take n $ m <| l
      yield l'
      go l'

--------------------------------------------------------------------------------
-- Random helpers

replaceTMVar :: TMVar a -> a -> STM ()
replaceTMVar var v = tryTakeTMVar var >> putTMVar var v


execU :: Monad m => m () -> Pipe a a m r
execU mOp = forever $ do
  lift mOp
  await >>= yield

--------------------------------------------------------------------------------
-- Status Icon handling

data MeasuredIcon
  = MI { miIcon :: Gtk.StatusIcon
       , miPixbuf :: TMVar Gtk.Pixbuf
       }

type Pixel = Word32


createIcon :: IO MeasuredIcon
createIcon = do
  icon <- Gtk.statusIconNew
  pixbufVar <- newEmptyTMVarIO
  _ <- Gtk.on icon Gtk.statusIconSizeChanged $ \size -> do
    pixbuf <- Gtk.pixbufNew Gtk.ColorspaceRgb True 8 size size
    atomically $ replaceTMVar pixbufVar pixbuf
    return True

  return $ MI icon pixbufVar

updateIcon :: MeasuredIcon -> Pixel -> Pixel -> Consumer (Seq Measurement) IO r
updateIcon (MI icon pixbufVar) color1 color2 = forever $ do
  measurements <- await
  pixbuf <- lift $ atomically $ readTMVar pixbufVar
  n <- lift $ Gtk.pixbufGetWidth pixbuf
  lift $ drawPixbuf pixbuf (S.take n measurements) color1 color2
  lift $ Gtk.postGUIAsync $ Gtk.statusIconSetFromPixbuf icon pixbuf

drawPixbuf :: (FoldableWithIndex Int t) => Gtk.Pixbuf -> t Measurement
              -> Pixel -> Pixel -> IO ()
drawPixbuf pb measurements color1 color2 = do
  Clear the pixbuf first
  Gtk.pixbufFill pb 0 0 0 0
  h <- Gtk.pixbufGetHeight pb
  w <- Gtk.pixbufGetWidth pb
  rs <- (`div` 4) <$> Gtk.pixbufGetRowstride pb
  pixels <- Gtk.pixbufGetPixels pb :: IO (Gtk.PixbufData Int Pixel)
  TODO(klao ): do this much more efficiently
  let
    setP x y c = MA.writeArray pixels ((h-1-y)*rs + (w-1-x)) c
    drawC x y0 y1 c = forM_ [y0..y1-1] $ \y -> setP x y c
    toH m = max 0 $ min h $ round $ fromIntegral h * m
    drawM x (m1,m2) = do
      drawC x 0 y1 color1
      drawC x y1 y2 color2
        where
          y1 = toH m1
          y2 = toH (m1+m2)

  imapMOf_ ifolded drawM measurements


main :: IO ()
main = do
  nonGtkArgs <- Gtk.unsafeInitGUIForThreadedRTS
  [] <- withArgs nonGtkArgs $ $initHFlags "Show CPU and MEM graph with tray icons"

  when flags_show_cpu $ do
    -- Colors are in ABGR (RGBA little-endian)
    let blue1 = 0xffff4400
        blue2 = 0x88ff4400
    cpuIcon <- createIcon
    _ <- forkIO $ runEffect $ measureCPU >-> delay >-> collectN 100
                            >-> updateIcon cpuIcon blue1 blue2
    return ()

  when flags_show_mem $ do
    let green1 = 0xff44ff00
        green2 = 0x8844ff00
    memIcon <- createIcon
    _ <- forkIO $ runEffect $ measureMEM >-> delay >-> collectN 100
                            >-> updateIcon memIcon green1 green2
    return ()

  -- Properly handle Ctrl-C:
  _ <- installHandler sigINT (Catch $ Gtk.postGUIAsync Gtk.mainQuit) Nothing
  Gtk.mainGUI


 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     -2.0

 Unless required by applicable law or agreed to in writing, software
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
# LANGUAGE OverloadedStrings #
------------------------------------------------------------------------------
# UNPACK #
# UNPACK #
# UNPACK #
 The counters are in milliseconds and cumulative since system start.
# UNPACK #
# UNPACK #
# UNPACK #
 If you stat /proc/stat or /proc/meminfo they look 0 sized, so we
 cannot use BS.readFile.
------------------------------------------------------------------------------
 Pipes
------------------------------------------------------------------------------
 Random helpers
------------------------------------------------------------------------------
 Status Icon handling
 Colors are in ABGR (RGBA little-endian)
 Properly handle Ctrl-C:
  Copyright 2014 Mihaly Barasz & Gergely Risko
  Licensed under the Apache License , Version 2.0 ( the " License " ) ;
  distributed under the License is distributed on an " AS IS " BASIS ,


# LANGUAGE TemplateHaskell #
# LANGUAGE FlexibleContexts #

module Main (main) where

import Control.Applicative
import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.STM
import Control.Lens
import Control.Monad
import qualified Data.Array.MArray as MA
import Data.Attoparsec.ByteString.Char8 (decimal, (.*>), skipSpace, skipWhile)
import qualified Data.Attoparsec.ByteString as AP
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Char
import Data.Sequence (Seq)
import qualified Data.Sequence as S
import Data.Word (Word32)
import qualified Graphics.UI.Gtk as Gtk
import HFlags
import Pipes (Producer, yield, (>->), await, Consumer, Pipe, lift, runEffect)
import System.Environment
import System.IO
import System.Posix.Signals
import Text.Printf


defineFlag "show_cpu" True "Show the CPU load indicator/icon."
defineFlag "show_mem" True "Show the memory usage indicator/icon."
defineFlag "update_interval" (500 :: Int)
  "Update interval for all indicators in milliseconds."

$(return [])

  Acquiring raw CPU and RAM numbers

data CPUCounters =
              } deriving (Show)

readCPUCounters :: IO CPUCounters
readCPUCounters = parseProcStat <$> readProcFile "/proc/stat"
  where
    parseProcStat = forceParse parseFirstLine
    parseFirstLine = wrapResult <$ "cpu " .*> skipSpace
                     <*> decimal <* skipSpace <*> decimal <* skipSpace
                     <*> decimal <* skipSpace <*> decimal <* skipSpace
                     <*> decimal <* skipSpace <*> decimal <* skipSpace
                     <*> decimal
    wrapResult user nice system idle iowait irq softirq =
      CPUCounters (user+nice) (system+iowait+irq+softirq) idle

data MEMCounters =
              } deriving (Show)

  The counters are in kB.
readMEMCounters :: IO MEMCounters
readMEMCounters = parseProcMeminfo <$> readProcFile "/proc/meminfo"
  where
    parseProcMeminfo = forceParse parseHead
    parseHead :: AP.Parser MEMCounters
    parseHead =
      wrapResult <$ "MemTotal:" .*> skipSpace <*> decimal <* kb <*
        "MemFree:" .*> skipSpace <*> decimal <* kb <*>
        maybeMemAvailable <*
        "Buffers:" .*> skipSpace <* (skipWhile isDigit) <* kb <*
        "Cached:" .*> skipSpace <*> decimal <* kb
    kb = skipSpace *> "kB" .*> skipSpace
    maybeMemAvailable =
      Just <$ "MemAvailable:" .*> skipSpace <*> decimal <* kb <|> pure Nothing
    wrapResult total free mAvailable cached =
      MEMCounters total cached $ maybe free id mAvailable

forceParse :: AP.Parser a -> BS.ByteString -> a
forceParse parser = either error id . AP.parseOnly parser

  We assume that all files of interest fit in 4k .
readProcFile :: FilePath -> IO ByteString
readProcFile file = withBinaryFile file ReadMode $ flip BS.hGetSome 4096


  TODO(klao ): use strict pairs
type Measurement = (Double, Double)

measureMEM :: Producer Measurement IO r
measureMEM = forever $ do
  MEMCounters total cached free <- lift readMEMCounters
  yield (fromIntegral (total-cached-free) / fromIntegral total,
         fromIntegral cached / fromIntegral total)

measureCPU :: Producer Measurement IO r
measureCPU = lift readCPUCounters >>= go
  where
    go (CPUCounters user0 sys0 idle0) = do
      m@(CPUCounters user sys idle) <- lift readCPUCounters
  TODO(klao ): handle overflows !
      let total = fromIntegral $ user + sys + idle - user0 - sys0 - idle0
      when (total > 0) $ yield (fromIntegral (user - user0) / total,
                                fromIntegral (sys - sys0) / total)
      go m

delay :: Pipe a a IO r
delay = execU (threadDelay $ flags_update_interval * 1000)

prettyPrint :: String -> Consumer Measurement IO r
prettyPrint label = forever $ do
  (x1, x2) <- await
  lift $ printf "%s: %4.1f%% %4.1f%%\n" label (x1*100) (x2*100)

collectN :: Int -> Pipe Measurement (Seq Measurement) IO r
collectN n = go S.empty
  where
    go l = do
      m <- await
      let l' = S.take n $ m <| l
      yield l'
      go l'


replaceTMVar :: TMVar a -> a -> STM ()
replaceTMVar var v = tryTakeTMVar var >> putTMVar var v


execU :: Monad m => m () -> Pipe a a m r
execU mOp = forever $ do
  lift mOp
  await >>= yield


data MeasuredIcon
  = MI { miIcon :: Gtk.StatusIcon
       , miPixbuf :: TMVar Gtk.Pixbuf
       }

type Pixel = Word32


createIcon :: IO MeasuredIcon
createIcon = do
  icon <- Gtk.statusIconNew
  pixbufVar <- newEmptyTMVarIO
  _ <- Gtk.on icon Gtk.statusIconSizeChanged $ \size -> do
    pixbuf <- Gtk.pixbufNew Gtk.ColorspaceRgb True 8 size size
    atomically $ replaceTMVar pixbufVar pixbuf
    return True

  return $ MI icon pixbufVar

updateIcon :: MeasuredIcon -> Pixel -> Pixel -> Consumer (Seq Measurement) IO r
updateIcon (MI icon pixbufVar) color1 color2 = forever $ do
  measurements <- await
  pixbuf <- lift $ atomically $ readTMVar pixbufVar
  n <- lift $ Gtk.pixbufGetWidth pixbuf
  lift $ drawPixbuf pixbuf (S.take n measurements) color1 color2
  lift $ Gtk.postGUIAsync $ Gtk.statusIconSetFromPixbuf icon pixbuf

drawPixbuf :: (FoldableWithIndex Int t) => Gtk.Pixbuf -> t Measurement
              -> Pixel -> Pixel -> IO ()
drawPixbuf pb measurements color1 color2 = do
  Clear the pixbuf first
  Gtk.pixbufFill pb 0 0 0 0
  h <- Gtk.pixbufGetHeight pb
  w <- Gtk.pixbufGetWidth pb
  rs <- (`div` 4) <$> Gtk.pixbufGetRowstride pb
  pixels <- Gtk.pixbufGetPixels pb :: IO (Gtk.PixbufData Int Pixel)
  TODO(klao ): do this much more efficiently
  let
    setP x y c = MA.writeArray pixels ((h-1-y)*rs + (w-1-x)) c
    drawC x y0 y1 c = forM_ [y0..y1-1] $ \y -> setP x y c
    toH m = max 0 $ min h $ round $ fromIntegral h * m
    drawM x (m1,m2) = do
      drawC x 0 y1 color1
      drawC x y1 y2 color2
        where
          y1 = toH m1
          y2 = toH (m1+m2)

  imapMOf_ ifolded drawM measurements


main :: IO ()
main = do
  nonGtkArgs <- Gtk.unsafeInitGUIForThreadedRTS
  [] <- withArgs nonGtkArgs $ $initHFlags "Show CPU and MEM graph with tray icons"

  when flags_show_cpu $ do
    let blue1 = 0xffff4400
        blue2 = 0x88ff4400
    cpuIcon <- createIcon
    _ <- forkIO $ runEffect $ measureCPU >-> delay >-> collectN 100
                            >-> updateIcon cpuIcon blue1 blue2
    return ()

  when flags_show_mem $ do
    let green1 = 0xff44ff00
        green2 = 0x8844ff00
    memIcon <- createIcon
    _ <- forkIO $ runEffect $ measureMEM >-> delay >-> collectN 100
                            >-> updateIcon memIcon green1 green2
    return ()

  _ <- installHandler sigINT (Catch $ Gtk.postGUIAsync Gtk.mainQuit) Nothing
  Gtk.mainGUI
