module Prometheus.Metric.Vector (
    Vector (..)
,   vector
,   withLabel
,   removeLabel
,   clearLabels
,   getVectorWith
) where

import Prometheus.Label
import Prometheus.Metric
import Prometheus.MonadMonitor

import Control.Applicative ((<$>))
import Control.DeepSeq
import qualified Data.Atomics as Atomics
import qualified Data.IORef as IORef
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import Data.Traversable (forM)


type VectorState l m = (Metric m, Map.Map l (m, IO [SampleGroup]))

data Vector l m = MkVector (IORef.IORef (VectorState l m))

instance NFData (Vector l m) where
  rnf :: Vector l m -> ()
rnf (MkVector IORef (VectorState l m)
ioref) = seq :: forall a b. a -> b -> b
seq IORef (VectorState l m)
ioref ()

-- | Creates a new vector of metrics given a label.
vector :: Label l => l -> Metric m -> Metric (Vector l m)
vector :: forall l m. Label l => l -> Metric m -> Metric (Vector l m)
vector l
labels Metric m
gen = forall s. IO (s, IO [SampleGroup]) -> Metric s
Metric forall a b. (a -> b) -> a -> b
$ do
    IORef (Metric m, Map l (m, IO [SampleGroup]))
ioref <- forall l a. Label l => l -> a -> a
checkLabelKeys l
labels forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
IORef.newIORef (Metric m
gen, forall k a. Map k a
Map.empty)
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall l m. IORef (VectorState l m) -> Vector l m
MkVector IORef (Metric m, Map l (m, IO [SampleGroup]))
ioref, forall l m.
Label l =>
l -> IORef (VectorState l m) -> IO [SampleGroup]
collectVector l
labels IORef (Metric m, Map l (m, IO [SampleGroup]))
ioref)

checkLabelKeys :: Label l => l -> a -> a
checkLabelKeys :: forall l a. Label l => l -> a -> a
checkLabelKeys l
keys a
r = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall {a}. a -> String -> a
check a
r forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Text -> String
T.unpack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall a b. (a -> b) -> a -> b
$ forall l. Label l => l -> l -> LabelPairs
labelPairs l
keys l
keys
    where
        check :: a -> String -> a
check a
_ String
"instance" = forall a. HasCallStack => String -> a
error String
"The label 'instance' is reserved."
        check a
_ String
"job"      = forall a. HasCallStack => String -> a
error String
"The label 'job' is reserved."
        check a
_ String
"quantile" = forall a. HasCallStack => String -> a
error String
"The label 'quantile' is reserved."
        check a
a (Char
k:String
ey)
            | Char -> Bool
validStart Char
k Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
validRest String
ey = a
a
            | Bool
otherwise = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"The label '" forall a. [a] -> [a] -> [a]
++ (Char
kforall a. a -> [a] -> [a]
:String
ey) forall a. [a] -> [a] -> [a]
++ String
"' is not valid."
        check a
_ []         = forall a. HasCallStack => String -> a
error String
"Empty labels are not allowed."

        validStart :: Char -> Bool
validStart Char
c =  (Char
'a' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'z')
                     Bool -> Bool -> Bool
|| (Char
'A' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'Z')
                     Bool -> Bool -> Bool
|| Char
c forall a. Eq a => a -> a -> Bool
== Char
'_'

        validRest :: Char -> Bool
validRest Char
c =  (Char
'a' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'z')
                    Bool -> Bool -> Bool
|| (Char
'A' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'Z')
                    Bool -> Bool -> Bool
|| (Char
'0' forall a. Ord a => a -> a -> Bool
<= Char
c Bool -> Bool -> Bool
&& Char
c forall a. Ord a => a -> a -> Bool
<= Char
'9')
                    Bool -> Bool -> Bool
|| Char
c forall a. Eq a => a -> a -> Bool
== Char
'_'

-- TODO(will): This currently makes the assumption that all the types and info
-- for all sample groups returned by a metric's collect method will be the same.
-- It is not clear that this will always be a valid assumption.
collectVector :: Label l => l -> IORef.IORef (VectorState l m) -> IO [SampleGroup]
collectVector :: forall l m.
Label l =>
l -> IORef (VectorState l m) -> IO [SampleGroup]
collectVector l
keys IORef (VectorState l m)
ioref = do
    (Metric m
_, Map l (m, IO [SampleGroup])
metricMap) <- forall a. IORef a -> IO a
IORef.readIORef IORef (VectorState l m)
ioref
    [SampleGroup] -> [SampleGroup]
joinSamples forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM forall {f :: * -> *} {a}.
Functor f =>
(l, (a, f [SampleGroup])) -> f [SampleGroup]
collectInner (forall k a. Map k a -> [(k, a)]
Map.assocs Map l (m, IO [SampleGroup])
metricMap)
    where
        collectInner :: (l, (a, f [SampleGroup])) -> f [SampleGroup]
collectInner (l
labels, (a
_metric, f [SampleGroup]
sampleGroups)) =
            forall a b. (a -> b) -> [a] -> [b]
map (l -> SampleGroup -> SampleGroup
adjustSamples l
labels) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f [SampleGroup]
sampleGroups

        adjustSamples :: l -> SampleGroup -> SampleGroup
adjustSamples l
labels (SampleGroup Info
info SampleType
ty [Sample]
samples) =
            Info -> SampleType -> [Sample] -> SampleGroup
SampleGroup Info
info SampleType
ty (forall a b. (a -> b) -> [a] -> [b]
map (l -> Sample -> Sample
prependLabels l
labels) [Sample]
samples)

        prependLabels :: l -> Sample -> Sample
prependLabels l
l (Sample Text
name LabelPairs
labels ByteString
value) =
            Text -> LabelPairs -> ByteString -> Sample
Sample Text
name (forall l. Label l => l -> l -> LabelPairs
labelPairs l
keys l
l forall a. [a] -> [a] -> [a]
++ LabelPairs
labels) ByteString
value

        joinSamples :: [SampleGroup] -> [SampleGroup]
joinSamples []                      = []
        joinSamples s :: [SampleGroup]
s@(SampleGroup Info
i SampleType
t [Sample]
_:[SampleGroup]
_) = [Info -> SampleType -> [Sample] -> SampleGroup
SampleGroup Info
i SampleType
t ([SampleGroup] -> [Sample]
extract [SampleGroup]
s)]

        extract :: [SampleGroup] -> [Sample]
extract [] = []
        extract (SampleGroup Info
_ SampleType
_ [Sample]
s:[SampleGroup]
xs) = [Sample]
s forall a. [a] -> [a] -> [a]
++ [SampleGroup] -> [Sample]
extract [SampleGroup]
xs

getVectorWith :: Vector label metric
              -> (metric -> IO a)
              -> IO [(label, a)]
getVectorWith :: forall label metric a.
Vector label metric -> (metric -> IO a) -> IO [(label, a)]
getVectorWith (MkVector IORef (VectorState label metric)
valueTVar) metric -> IO a
f = do
    (Metric metric
_, Map label (metric, IO [SampleGroup])
metricMap) <- forall a. IORef a -> IO a
IORef.readIORef IORef (VectorState label metric)
valueTVar
    forall k a. Map k a -> [(k, a)]
Map.assocs forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Map label (metric, IO [SampleGroup])
metricMap (metric -> IO a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst)

-- | Given a label, applies an operation to the corresponding metric in the
-- vector.
withLabel :: (Label label, MonadMonitor m)
          => Vector label metric
          -> label
          -> (metric -> IO ())
          -> m ()
withLabel :: forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
withLabel (MkVector IORef (VectorState label metric)
ioref) label
label metric -> IO ()
f = forall (m :: * -> *). MonadMonitor m => IO () -> m ()
doIO forall a b. (a -> b) -> a -> b
$ do
    (Metric IO (metric, IO [SampleGroup])
gen, Map label (metric, IO [SampleGroup])
_) <- forall a. IORef a -> IO a
IORef.readIORef IORef (VectorState label metric)
ioref
    (metric, IO [SampleGroup])
newMetric <- IO (metric, IO [SampleGroup])
gen
    (metric, IO [SampleGroup])
metric <- forall a b. IORef a -> (a -> (a, b)) -> IO b
Atomics.atomicModifyIORefCAS IORef (VectorState label metric)
ioref forall a b. (a -> b) -> a -> b
$ \(Metric metric
_, Map label (metric, IO [SampleGroup])
metricMap) ->
        let maybeMetric :: Maybe (metric, IO [SampleGroup])
maybeMetric = forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup label
label Map label (metric, IO [SampleGroup])
metricMap
            updatedMap :: Map label (metric, IO [SampleGroup])
updatedMap  = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert label
label (metric, IO [SampleGroup])
newMetric Map label (metric, IO [SampleGroup])
metricMap
        in  case Maybe (metric, IO [SampleGroup])
maybeMetric of
                Maybe (metric, IO [SampleGroup])
Nothing     -> ((forall s. IO (s, IO [SampleGroup]) -> Metric s
Metric IO (metric, IO [SampleGroup])
gen, Map label (metric, IO [SampleGroup])
updatedMap), (metric, IO [SampleGroup])
newMetric)
                Just (metric, IO [SampleGroup])
metric -> ((forall s. IO (s, IO [SampleGroup]) -> Metric s
Metric IO (metric, IO [SampleGroup])
gen, Map label (metric, IO [SampleGroup])
metricMap), (metric, IO [SampleGroup])
metric)
    metric -> IO ()
f (forall a b. (a, b) -> a
fst (metric, IO [SampleGroup])
metric)

-- | Removes a label from a vector.
removeLabel :: (Label label, MonadMonitor m)
            => Vector label metric -> label -> m ()
removeLabel :: forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> m ()
removeLabel (MkVector IORef (VectorState label metric)
valueTVar) label
label =
    forall (m :: * -> *). MonadMonitor m => IO () -> m ()
doIO forall a b. (a -> b) -> a -> b
$ forall t. IORef t -> (t -> t) -> IO ()
Atomics.atomicModifyIORefCAS_ IORef (VectorState label metric)
valueTVar forall {a} {a}. (a, Map label a) -> (a, Map label a)
f
    where f :: (a, Map label a) -> (a, Map label a)
f (a
desc, Map label a
metricMap) = (a
desc, forall k a. Ord k => k -> Map k a -> Map k a
Map.delete label
label Map label a
metricMap)

-- | Removes all labels from a vector.
clearLabels :: (Label label, MonadMonitor m)
            => Vector label metric -> m ()
clearLabels :: forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> m ()
clearLabels (MkVector IORef (VectorState label metric)
valueTVar) =
    forall (m :: * -> *). MonadMonitor m => IO () -> m ()
doIO forall a b. (a -> b) -> a -> b
$ forall t. IORef t -> (t -> t) -> IO ()
Atomics.atomicModifyIORefCAS_ IORef (VectorState label metric)
valueTVar forall {a} {b} {k} {a}. (a, b) -> (a, Map k a)
f
    where f :: (a, b) -> (a, Map k a)
f (a
desc, b
_) = (a
desc, forall k a. Map k a
Map.empty)