{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Database.Redis.ManualCommands.Cms where

import Data.ByteString (ByteString)
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NE
import Data.Maybe (listToMaybe)

import Database.Redis.Core
import Database.Redis.Protocol
import Database.Redis.Types

data CMSInfo = CMSInfo
    { CMSInfo -> Integer
cmsInfoWidth :: Integer
      -- ^ Number of counters in each array row.
    , CMSInfo -> Integer
cmsInfoDepth :: Integer
      -- ^ Number of counter array rows.
    , CMSInfo -> Integer
cmsInfoCount :: Integer
      -- ^ Total count added to the sketch.
    } deriving (Int -> CMSInfo -> ShowS
[CMSInfo] -> ShowS
CMSInfo -> String
(Int -> CMSInfo -> ShowS)
-> (CMSInfo -> String) -> ([CMSInfo] -> ShowS) -> Show CMSInfo
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CMSInfo -> ShowS
showsPrec :: Int -> CMSInfo -> ShowS
$cshow :: CMSInfo -> String
show :: CMSInfo -> String
$cshowList :: [CMSInfo] -> ShowS
showList :: [CMSInfo] -> ShowS
Show, CMSInfo -> CMSInfo -> Bool
(CMSInfo -> CMSInfo -> Bool)
-> (CMSInfo -> CMSInfo -> Bool) -> Eq CMSInfo
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CMSInfo -> CMSInfo -> Bool
== :: CMSInfo -> CMSInfo -> Bool
$c/= :: CMSInfo -> CMSInfo -> Bool
/= :: CMSInfo -> CMSInfo -> Bool
Eq)

instance RedisResult CMSInfo where
    decode :: Reply -> Either Reply CMSInfo
decode Reply
r = do
        fields <- Reply -> Either Reply [(ByteString, Integer)]
forall a. RedisResult a => Reply -> Either Reply a
decode Reply
r :: Either Reply [(ByteString, Integer)]
        cmsInfoWidth <- decodeField ["width", "Width"] fields
        cmsInfoDepth <- decodeField ["depth", "Depth"] fields
        cmsInfoCount <- decodeField ["count", "Count"] fields
        pure CMSInfo{..}
      where
        decodeField :: [a] -> [(a, b)] -> Either Reply b
decodeField [a]
keys [(a, b)]
fields =
            Either Reply b
-> (b -> Either Reply b) -> Maybe b -> Either Reply b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Reply -> Either Reply b
forall a b. a -> Either a b
Left Reply
r) b -> Either Reply b
forall a b. b -> Either a b
Right (Maybe b -> Either Reply b)
-> ([b] -> Maybe b) -> [b] -> Either Reply b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [b] -> Maybe b
forall a. [a] -> Maybe a
listToMaybe ([b] -> Either Reply b) -> [b] -> Either Reply b
forall a b. (a -> b) -> a -> b
$
                [b
value | a
key <- [a]
keys, b
value <- Maybe b -> [b]
forall {a}. Maybe a -> [a]
maybeToList (a -> [(a, b)] -> Maybe b
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup a
key [(a, b)]
fields)]

        maybeToList :: Maybe a -> [a]
maybeToList = [a] -> (a -> [a]) -> Maybe a -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] a -> [a]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure

data CMSMergeOpts
    = CMSMergeUnweighted (NonEmpty ByteString)
      -- ^ Merge the given source sketches using the default weight of @1@.
    | CMSMergeWeighted (NonEmpty (ByteString, Integer))
      -- ^ Merge the given source sketches using an explicit weight for each source.
    deriving (Int -> CMSMergeOpts -> ShowS
[CMSMergeOpts] -> ShowS
CMSMergeOpts -> String
(Int -> CMSMergeOpts -> ShowS)
-> (CMSMergeOpts -> String)
-> ([CMSMergeOpts] -> ShowS)
-> Show CMSMergeOpts
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CMSMergeOpts -> ShowS
showsPrec :: Int -> CMSMergeOpts -> ShowS
$cshow :: CMSMergeOpts -> String
show :: CMSMergeOpts -> String
$cshowList :: [CMSMergeOpts] -> ShowS
showList :: [CMSMergeOpts] -> ShowS
Show, CMSMergeOpts -> CMSMergeOpts -> Bool
(CMSMergeOpts -> CMSMergeOpts -> Bool)
-> (CMSMergeOpts -> CMSMergeOpts -> Bool) -> Eq CMSMergeOpts
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CMSMergeOpts -> CMSMergeOpts -> Bool
== :: CMSMergeOpts -> CMSMergeOpts -> Bool
$c/= :: CMSMergeOpts -> CMSMergeOpts -> Bool
/= :: CMSMergeOpts -> CMSMergeOpts -> Bool
Eq)

cmsMergeOptsToArgs :: CMSMergeOpts -> [ByteString]
cmsMergeOptsToArgs :: CMSMergeOpts -> [ByteString]
cmsMergeOptsToArgs (CMSMergeUnweighted NonEmpty ByteString
sourceKeys) =
    Integer -> ByteString
forall a. RedisArg a => a -> ByteString
encode (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (NonEmpty ByteString -> Int
forall a. NonEmpty a -> Int
NE.length NonEmpty ByteString
sourceKeys) :: Integer) ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: NonEmpty ByteString -> [ByteString]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty ByteString
sourceKeys
cmsMergeOptsToArgs (CMSMergeWeighted NonEmpty (ByteString, Integer)
weightedSourceKeys) =
    Integer -> ByteString
forall a. RedisArg a => a -> ByteString
encode (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (NonEmpty (ByteString, Integer) -> Int
forall a. NonEmpty a -> Int
NE.length NonEmpty (ByteString, Integer)
weightedSourceKeys) :: Integer)
        ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: ((ByteString, Integer) -> ByteString)
-> [(ByteString, Integer)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (ByteString, Integer) -> ByteString
forall a b. (a, b) -> a
fst [(ByteString, Integer)]
sourceKeyWeights
        [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString
"WEIGHTS"]
        [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ ((ByteString, Integer) -> ByteString)
-> [(ByteString, Integer)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> ByteString
forall a. RedisArg a => a -> ByteString
encode (Integer -> ByteString)
-> ((ByteString, Integer) -> Integer)
-> (ByteString, Integer)
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString, Integer) -> Integer
forall a b. (a, b) -> b
snd) [(ByteString, Integer)]
sourceKeyWeights
  where
    sourceKeyWeights :: [(ByteString, Integer)]
sourceKeyWeights = NonEmpty (ByteString, Integer) -> [(ByteString, Integer)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (ByteString, Integer)
weightedSourceKeys

-- |Increases the count of one or more items by increment (<https://redis.io/commands/cms.incrby>).
--
-- /O(n)/ where /n/ is the number of items
--
-- Since RedisBloom 2.0.0
cmsincrby
    :: (RedisCtx m f)
    => ByteString -- ^ Key of the Count-Min Sketch.
    -> NonEmpty (ByteString, Integer) -- ^ Item and increment pairs.
    -> m (f [Integer])
cmsincrby :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> NonEmpty (ByteString, Integer) -> m (f [Integer])
cmsincrby ByteString
key NonEmpty (ByteString, Integer)
itemIncrements =
    [ByteString] -> m (f [Integer])
forall (m :: * -> *) (f :: * -> *) a.
(RedisCtx m f, RedisResult a) =>
[ByteString] -> m (f a)
sendRequest ([ByteString] -> m (f [Integer]))
-> [ByteString] -> m (f [Integer])
forall a b. (a -> b) -> a -> b
$ [ByteString
"CMS.INCRBY", ByteString
key] [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ ((ByteString, Integer) -> [ByteString])
-> [(ByteString, Integer)] -> [ByteString]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (ByteString, Integer) -> [ByteString]
forall {a}. RedisArg a => (ByteString, a) -> [ByteString]
encodeItemIncrement (NonEmpty (ByteString, Integer) -> [(ByteString, Integer)]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty (ByteString, Integer)
itemIncrements)
  where
    encodeItemIncrement :: (ByteString, a) -> [ByteString]
encodeItemIncrement (ByteString
item, a
increment) = [ByteString
item, a -> ByteString
forall a. RedisArg a => a -> ByteString
encode a
increment]

-- |Returns information about a sketch (<https://redis.io/commands/cms.info>).
--
-- /O(1)/
--
-- Since RedisBloom 2.0.0
cmsinfo
    :: (RedisCtx m f)
    => ByteString -- ^ Key of the Count-Min Sketch.
    -> m (f CMSInfo)
cmsinfo :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> m (f CMSInfo)
cmsinfo ByteString
key = [ByteString] -> m (f CMSInfo)
forall (m :: * -> *) (f :: * -> *) a.
(RedisCtx m f, RedisResult a) =>
[ByteString] -> m (f a)
sendRequest [ByteString
"CMS.INFO", ByteString
key]

-- |Initializes a Count-Min Sketch to dimensions specified by user (<https://redis.io/commands/cms.initbydim>).
--
-- /O(1)/
--
-- Since RedisBloom 2.0.0
cmsinitbydim
    :: (RedisCtx m f)
    => ByteString -- ^ Key of the Count-Min Sketch to create.
    -> Integer -- ^ Number of counters in each row.
    -> Integer -- ^ Number of counter rows.
    -> m (f Status)
cmsinitbydim :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> Integer -> Integer -> m (f Status)
cmsinitbydim ByteString
key Integer
width Integer
depth =
    [ByteString] -> m (f Status)
forall (m :: * -> *) (f :: * -> *) a.
(RedisCtx m f, RedisResult a) =>
[ByteString] -> m (f a)
sendRequest [ByteString
"CMS.INITBYDIM", ByteString
key, Integer -> ByteString
forall a. RedisArg a => a -> ByteString
encode Integer
width, Integer -> ByteString
forall a. RedisArg a => a -> ByteString
encode Integer
depth]

-- |Initializes a Count-Min Sketch to accommodate requested tolerances (<https://redis.io/commands/cms.initbyprob>).
--
-- /O(1)/
--
-- Since RedisBloom 2.0.0
cmsinitbyprob
    :: (RedisCtx m f)
    => ByteString -- ^ Key of the Count-Min Sketch to create.
    -> Double -- ^ Error factor.
    -> Double -- ^ Probability of the error factor.
    -> m (f Status)
cmsinitbyprob :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> Double -> Double -> m (f Status)
cmsinitbyprob ByteString
key Double
err Double
probability =
    [ByteString] -> m (f Status)
forall (m :: * -> *) (f :: * -> *) a.
(RedisCtx m f, RedisResult a) =>
[ByteString] -> m (f a)
sendRequest [ByteString
"CMS.INITBYPROB", ByteString
key, Double -> ByteString
forall a. RedisArg a => a -> ByteString
encode Double
err, Double -> ByteString
forall a. RedisArg a => a -> ByteString
encode Double
probability]

-- |Merges several sketches into one sketch (<https://redis.io/commands/cms.merge>).
--
-- Source sketches are merged with the default weight of @1@.
--
-- /O(n)/ where /n/ is the number of sketches
--
-- Since RedisBloom 2.0.0
cmsmerge
    :: (RedisCtx m f)
    => ByteString -- ^ Destination sketch key.
    -> NonEmpty ByteString -- ^ Source sketch keys.
    -> m (f Status)
cmsmerge :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> NonEmpty ByteString -> m (f Status)
cmsmerge ByteString
destination NonEmpty ByteString
sourceKeys =
    ByteString -> CMSMergeOpts -> m (f Status)
forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> CMSMergeOpts -> m (f Status)
cmsmergeOpts ByteString
destination (NonEmpty ByteString -> CMSMergeOpts
CMSMergeUnweighted NonEmpty ByteString
sourceKeys)

-- |Merges several sketches into one sketch (<https://redis.io/commands/cms.merge>).
--
-- /O(n)/ where /n/ is the number of sketches
--
-- Since RedisBloom 2.0.0
cmsmergeOpts
    :: (RedisCtx m f)
    => ByteString -- ^ Destination sketch key.
    -> CMSMergeOpts -- ^ Source sketches with optional weights.
    -> m (f Status)
cmsmergeOpts :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> CMSMergeOpts -> m (f Status)
cmsmergeOpts ByteString
destination CMSMergeOpts
opts =
    [ByteString] -> m (f Status)
forall (m :: * -> *) (f :: * -> *) a.
(RedisCtx m f, RedisResult a) =>
[ByteString] -> m (f a)
sendRequest ([ByteString] -> m (f Status)) -> [ByteString] -> m (f Status)
forall a b. (a -> b) -> a -> b
$ [ByteString
"CMS.MERGE", ByteString
destination] [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ CMSMergeOpts -> [ByteString]
cmsMergeOptsToArgs CMSMergeOpts
opts

-- |Merges several sketches into one sketch (<https://redis.io/commands/cms.merge>).
--
-- /O(n)/ where /n/ is the number of sketches
--
-- Since RedisBloom 2.0.0
cmsmergeWeighted
    :: (RedisCtx m f)
    => ByteString -- ^ Destination sketch key.
    -> NonEmpty (ByteString, Integer) -- ^ Source sketch keys paired with weights.
    -> m (f Status)
cmsmergeWeighted :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> NonEmpty (ByteString, Integer) -> m (f Status)
cmsmergeWeighted ByteString
destination NonEmpty (ByteString, Integer)
weightedSourceKeys =
    ByteString -> CMSMergeOpts -> m (f Status)
forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> CMSMergeOpts -> m (f Status)
cmsmergeOpts ByteString
destination (NonEmpty (ByteString, Integer) -> CMSMergeOpts
CMSMergeWeighted NonEmpty (ByteString, Integer)
weightedSourceKeys)

-- |Returns the count for one or more items in a sketch (<https://redis.io/commands/cms.query>).
--
-- /O(n)/ where /n/ is the number of items
--
-- Since RedisBloom 2.0.0
cmsquery
    :: (RedisCtx m f)
    => ByteString -- ^ Key of the Count-Min Sketch.
    -> NonEmpty ByteString -- ^ Items to query.
    -> m (f [Integer])
cmsquery :: forall (m :: * -> *) (f :: * -> *).
RedisCtx m f =>
ByteString -> NonEmpty ByteString -> m (f [Integer])
cmsquery ByteString
key NonEmpty ByteString
items = [ByteString] -> m (f [Integer])
forall (m :: * -> *) (f :: * -> *) a.
(RedisCtx m f, RedisResult a) =>
[ByteString] -> m (f a)
sendRequest ([ByteString] -> m (f [Integer]))
-> [ByteString] -> m (f [Integer])
forall a b. (a -> b) -> a -> b
$ [ByteString
"CMS.QUERY", ByteString
key] [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ NonEmpty ByteString -> [ByteString]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty ByteString
items