{-# LANGUAGE LambdaCase, OverloadedStrings #-}
{- Some references:
 -  https://blog.chih.me/read-cpu-power-with-RAPL.html
 -  https://github.com/torvalds/linux/blob/master/Documentation/power/powercap/powercap.txt
 -  https://01.org/blogs/2014/running-average-power-limit-–-rapl
 -  http://web.eece.maine.edu/~vweaver/projects/rapl/
 -}

module System.Power.IntelRAPL (
    PowerDomain, Joule
  , PowerDomains
  , enumeratePowerDomains, showPowerDomains
  , usedEnergy, usedSince
  , showEnergy
  ) where

import           Data.Fixed
import           Data.Foldable
import           Data.List
import           Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Map.Merge.Lazy as Map
import           Data.Maybe (mapMaybe)
import           Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import           Data.Tree (Tree, unfoldTreeM, drawTree)
import qualified Data.Tree as Tree
import           System.Directory (listDirectory)
import           System.FilePath

type Joule = Micro

data PowerDomain =
    PowerDomain
    { _pdName :: Text
    , _pdDomain :: FilePath
    }
  deriving (Read, Show, Eq, Ord)

type PowerDomains = Tree (Either Text PowerDomain)

sysRaplPrefix :: FilePath
sysRaplPrefix = "/sys/class/powercap/intel-rapl:"

showPowerDomains :: PowerDomains -> String
showPowerDomains = drawTree . fmap (T.unpack . either id _pdName)

enumeratePowerDomains :: IO PowerDomains
enumeratePowerDomains = do
    domains <- (sort . mapMaybe ((takeFileName sysRaplPrefix) `stripPrefix`)) <$>
               listDirectory (takeDirectory sysRaplPrefix)
    unfoldTreeM expandDomain (Left "System", domains)
  where
    subTrees :: Eq a => [[a]] -> [(Either b [a], [[a]])]
    subTrees = unfoldr (\case
                           [] -> Nothing
                           h:t -> let (sd, o) = partition ((h `isPrefixOf`)) t
                                  in Just ((Right h, sd), o))
    expandDomain :: (Either Text FilePath, [FilePath])
                 -> IO (Either Text PowerDomain, [(Either Text FilePath, [FilePath])])
    expandDomain (Left nm, subDomains) =
      return (Left nm,  subTrees subDomains)
    expandDomain (Right dmn, subDomains) = do
      let dr =  sysRaplPrefix ++ dmn
      nm <- T.strip <$> TIO.readFile (dr </> "name")
      return (Right (PowerDomain nm dr), subTrees subDomains)

usedEnergy :: PowerDomains -> IO (Map PowerDomain Joule)
usedEnergy t =
  fmap Map.fromList . foldrM (\a b -> case a of
                    Left _ -> return b
                    Right pd@(PowerDomain _ d) ->
                      (\j -> (pd, MkFixed (read j)):b) <$> readFile (d</>"energy_uj")) [] $ t

usedSince :: Map PowerDomain Joule -> Map PowerDomain Joule -> Map PowerDomain Joule
usedSince = Map.merge Map.dropMissing Map.dropMissing (Map.zipWithMatched $ const (-))

showEnergy :: PowerDomains -> Map PowerDomain Joule -> String
showEnergy pds js =
    showETree 0 pds
  where
    showETree :: Int -> PowerDomains -> String
    showETree dp (Tree.Node (Left nm) sbs) =
      replicate dp ' '++"-"++T.unpack nm++": Unk J\n"++
      (concatMap (showETree (dp+1)) sbs)
    showETree dp (Tree.Node (Right (pd@(PowerDomain nm _))) sbs) =
      replicate dp ' '++"-"++T.unpack nm++": "++(maybe "Unk" show . Map.lookup pd $ js)++" J\n"++
      (concatMap (showETree (dp+1)) sbs)