{-# LANGUAGE OverloadedStrings #-}

module System.Metrics.RRDTool.Internals where

import Control.Monad
import Control.Monad.Writer
import Data.Int
import Data.List
import Data.Time
import Data.Ord
import Data.Word
import System.Metrics
import Text.Printf

import qualified Data.HashMap.Strict as HM
import qualified Data.Text as T

{- | Intervals in seconds (e.g. heartbeats and step size) -}
type IntervalSeconds = Int

{- | Counters, gauges etc. in `ekg` have this type -}
type SourceValue = Int64

{- | Types of data source in a round-robin database. -}
data DataSourceType
  {- | Metrics whose current value is tracked. This is appropriate for
'System.Metrics.Gauge.Gauge' metrics and creates a `GAUGE` DS. -}
  = DsGauge
    {- | Metrics that are monotonically increasing and whose rate-of-change is
tracked. This is appropriate for 'System.Metrics.Counter.Counter' metrics and
creates a `DERIVE` DS. -}
  | DsDerive
  deriving (Eq, Show, Bounded, Enum)

{- | A data source in a round-robin database. -}
data DataSource = DataSource
  { dsName      :: T.Text -- ^ The name of the metric in the database. Maximum 19 characters and contains letters, digits and underscores only.
  , dsMetric    :: T.Text -- ^ The name of the metric in the `ekg` 'Store'.
  , dsType      :: DataSourceType -- ^ The type of the metric.
  , dsHeartBeat :: IntervalSeconds -- ^ After this length of time with no updates, consider the value to be \'unknown\'.
  , dsMin       :: Maybe SourceValue -- ^ Values less than this should be interpreted as \'unknown\'.
  , dsMax       :: Maybe SourceValue -- ^ Values greater than this should be interpreted as \'unknown\'.
  } deriving (Show, Eq)

{- | Pre-defined data sources for tracking GHC's GC metrics. See 'registerGcMetrics' or 'GHC.Stats.GCStats' for more details.
     Clients must run 'registerGcMetrics' to register these metrics once the RRD is created. -}
gcSources
  :: IntervalSeconds -- ^ The heartbeat for these metrics, in seconds.
  -> [DataSource]
gcSources heartBeat = gcCounters ++ gcGauges
  where
  gcCounters = map (\(metric, name) -> DataSource name metric DsDerive heartBeat (Just 0) Nothing)
      [("rts.gc.bytes_allocated", "bytes_allocated")
      ,("rts.gc.num_gcs", "num_gcs")
      ,("rts.gc.num_bytes_usage_samples", "num_bytes_usage_sam")
      ,("rts.gc.cumulative_bytes_used", "cumulative_bytes_us")
      ,("rts.gc.bytes_copied", "bytes_copied")
      ,("rts.gc.mutator_cpu_ms", "mutator_cpu_ms")
      ,("rts.gc.mutator_wall_ms", "mutator_wall_ms")
      ,("rts.gc.gc_cpu_ms", "gc_cpu_ms")
      ,("rts.gc.gc_wall_ms", "gc_wall_ms")
      ,("rts.gc.cpu_ms", "cpu_ms")
      ,("rts.gc.wall_ms", "wall_ms")
      ]

  gcGauges = map (\(metric, name) -> DataSource name metric DsGauge heartBeat (Just 0) Nothing)
      [("rts.gc.max_bytes_used", "max_bytes_used")
      ,("rts.gc.current_bytes_used", "current_bytes_used")
      ,("rts.gc.current_bytes_slop", "current_bytes_slop")
      ,("rts.gc.max_bytes_slop", "max_bytes_slop")
      ,("rts.gc.peak_megabytes_allocated", "peak_megabytes_allo")
      ]

{- | Defines how multiple primary data points (PDPs) are turned into a
consolidated data point (CDP) for storage in a 'RoundRobinArchive'. -}
data ConsolidationFunction
  = CFLast -- ^ Keep the last point only.
  | CFAverage -- ^ Take the mean of the points.
  | CFMin -- ^ Take the minimum point.
  | CFMax -- ^ Take the maximum point.
  deriving (Show, Eq)

{- | A sequence of consolidated data points (CDPs). -}
data RoundRobinArchive = RoundRobinArchive
  { rraCf :: ConsolidationFunction -- ^ How to consolidate PDPs to get each CDP.
  , rraXff :: Double -- ^ The \'Xfiles factor\' - if more than this proportion of PDPs is unknown then the CDP is unknown.
  , rraPdpCount :: Int -- ^ The number of PDPs to use to calculate each CDP.
  , rraRecordCount :: Int -- ^ The number of CDPs to store in the archive.
  } deriving (Show, Eq)

{- | A round-robin database (RRD), which is a file on disk that stores
     time series data from a number of sources. -}
data RoundRobinDatabase = RoundRobinDatabase
  { rrdToolPath :: FilePath
  , rrdFilePath :: FilePath
  , rrdSources  :: HM.HashMap T.Text DataSource
  , rrdArchives :: [RoundRobinArchive]
  , rrdStore    :: Store -- ^ Get the 'Store' associated with this 'RoundRobinDatabase' in order to register metrics.
  , rrdStep     :: IntervalSeconds
  }

defineDataSource :: DataSource -> String
defineDataSource ds = printf "DS:%s:%s:%d:%s:%s"
  (T.unpack $ dsName ds)
  (case dsType ds of
    DsDerive -> "DERIVE"
    DsGauge  -> "GAUGE" :: String)
  (dsHeartBeat ds)
  (showMaybeValue $ dsMin ds)
  (showMaybeValue $ dsMax ds)

showMaybeValue :: Maybe SourceValue -> String
showMaybeValue = maybe "U" (showWord64 . fromIntegral)

defineRoundRobinArchive :: RoundRobinArchive -> String
defineRoundRobinArchive rra = printf "RRA:%s:%f:%d:%d"
  (case rraCf rra of
    CFAverage -> "AVERAGE"
    CFMin -> "MIN"
    CFMax -> "MAX"
    CFLast -> "LAST" :: String)
  (rraXff rra)
  (rraPdpCount rra)
  (rraRecordCount rra)

createRRDArgs :: RoundRobinDatabase -> [String]
createRRDArgs rrd = execWriter $ do
  tell ["create", rrdFilePath rrd, "--no-overwrite", "-s", show $ rrdStep rrd]
  forM_ (sortBy (comparing dsName) $ HM.elems $ rrdSources rrd)
    $ tell . return . defineDataSource
  forM_ (rrdArchives rrd) $ tell . return . defineRoundRobinArchive

updateRRDArgs :: RoundRobinDatabase -> HM.HashMap T.Text Value -> UTCTime -> Maybe [String]
updateRRDArgs rrd sample now =
  let updateSpec = HM.elems $ HM.intersectionWith (,) sample $ rrdSources rrd
  in if null updateSpec
        then Nothing
        else Just $ execWriter $ do
          tell ["update", rrdFilePath rrd]
          tell ["-t", intercalate ":" $ map (T.unpack . dsName . snd) updateSpec]
          tell [epochTimeFromUTCTime now ++ concatMap (formatValue . fst) updateSpec]

epoch :: UTCTime
epoch = UTCTime (fromGregorian 1970 1 1) 0

epochTime :: UTCTime -> Integer
epochTime = floor . flip diffUTCTime epoch

epochTimeFromUTCTime :: UTCTime -> String
epochTimeFromUTCTime = show . epochTime

formatValue :: Value -> String
formatValue = (':':) . showMaybeValue . getMaybeValue
  where getMaybeValue (Counter n) = Just n
        getMaybeValue (Gauge n)   = Just n
        getMaybeValue _           = Nothing

showWord64 :: Word64 -> String
showWord64 = show