{-# LANGUAGE TupleSections #-}

module System.Metrics.Prometheus.Metric.Histogram (
    Histogram,
    HistogramSample (..),
    Buckets,
    UpperBound,
    new,
    observe,
    sample,
    observeAndSample,
) where

import Control.Applicative ((<$>))
import Control.Monad (void)
import Data.Bool (bool)
import Data.IORef (
    IORef,
    atomicModifyIORef',
    newIORef,
    readIORef,
 )
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map


newtype Histogram = Histogram {unHistogram :: IORef HistogramSample}


type UpperBound = Double -- Inclusive upper bounds
type Buckets = Map UpperBound Double


data HistogramSample = HistogramSample
    { histBuckets :: !Buckets
    , histSum :: !Double
    , histCount :: !Int
    }


new :: [UpperBound] -> IO Histogram
new buckets = Histogram <$> newIORef empty
  where
    empty = HistogramSample (Map.fromList $ map (,0) (read "Infinity" : buckets)) zeroSum zeroCount
    zeroSum = 0.0
    zeroCount = 0


observeAndSample :: Double -> Histogram -> IO HistogramSample
observeAndSample x = flip atomicModifyIORef' update . unHistogram
  where
    update histData = (hist' histData, histData)
    hist' histData =
        histData
            { histBuckets = updateBuckets x $ histBuckets histData
            , histSum = histSum histData + x
            , histCount = histCount histData + 1
            }


observe :: Double -> Histogram -> IO ()
observe x = void . observeAndSample x


updateBuckets :: Double -> Buckets -> Buckets
updateBuckets x = Map.mapWithKey updateBucket
  where
    updateBucket key val = bool val (val + 1) (x <= key)


sample :: Histogram -> IO HistogramSample
sample = readIORef . unHistogram