    {-# LANGUAGE OverloadedStrings, RecordWildCards #-}
    module Database.Redis.Cluster.Command where

    import Data.Char(toLower)
    import qualified Data.ByteString as BS
    import qualified Data.ByteString.Char8 as Char8
    import qualified Data.HashMap.Strict as HM
    import Database.Redis.Types(RedisResult(decode))
    import Database.Redis.Protocol(Reply(..))

    data Flag
        = Write
        | ReadOnly
        | DenyOOM
        | Admin
        | PubSub
        | NoScript
        | Random
        | SortForScript
        | Loading
        | Stale
        | SkipMonitor
        | Asking
        | Fast
        | MovableKeys
        | Other BS.ByteString deriving (Int -> Flag -> ShowS
    [Flag] -> ShowS
    Flag -> [Char]
    (Int -> Flag -> ShowS)
    -> (Flag -> [Char]) -> ([Flag] -> ShowS) -> Show Flag
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Flag -> ShowS
    showsPrec :: Int -> Flag -> ShowS
    $cshow :: Flag -> [Char]
    show :: Flag -> [Char]
    $cshowList :: [Flag] -> ShowS
    showList :: [Flag] -> ShowS
    Show, Flag -> Flag -> Bool
    (Flag -> Flag -> Bool) -> (Flag -> Flag -> Bool) -> Eq Flag
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Flag -> Flag -> Bool
    == :: Flag -> Flag -> Bool
    $c/= :: Flag -> Flag -> Bool
    /= :: Flag -> Flag -> Bool
    Eq)

    data AritySpec = Required Integer | MinimumRequired Integer deriving (Int -> AritySpec -> ShowS
    [AritySpec] -> ShowS
    AritySpec -> [Char]
    (Int -> AritySpec -> ShowS)
    -> (AritySpec -> [Char])
    -> ([AritySpec] -> ShowS)
    -> Show AritySpec
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> AritySpec -> ShowS
    showsPrec :: Int -> AritySpec -> ShowS
    $cshow :: AritySpec -> [Char]
    show :: AritySpec -> [Char]
    $cshowList :: [AritySpec] -> ShowS
    showList :: [AritySpec] -> ShowS
    Show)

    data LastKeyPositionSpec = LastKeyPosition Integer | UnlimitedKeys Integer deriving (Int -> LastKeyPositionSpec -> ShowS
    [LastKeyPositionSpec] -> ShowS
    LastKeyPositionSpec -> [Char]
    (Int -> LastKeyPositionSpec -> ShowS)
    -> (LastKeyPositionSpec -> [Char])
    -> ([LastKeyPositionSpec] -> ShowS)
    -> Show LastKeyPositionSpec
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> LastKeyPositionSpec -> ShowS
    showsPrec :: Int -> LastKeyPositionSpec -> ShowS
    $cshow :: LastKeyPositionSpec -> [Char]
    show :: LastKeyPositionSpec -> [Char]
    $cshowList :: [LastKeyPositionSpec] -> ShowS
    showList :: [LastKeyPositionSpec] -> ShowS
    Show)

    newtype InfoMap = InfoMap (HM.HashMap String CommandInfo)

    -- Represents the result of the COMMAND command, which returns information
    -- about the position of keys in a request
    data CommandInfo = CommandInfo
        { CommandInfo -> ByteString
    name :: BS.ByteString
        , CommandInfo -> AritySpec
    arity :: AritySpec
        , CommandInfo -> [Flag]
    flags :: [Flag]
        , CommandInfo -> Integer
    firstKeyPosition :: Integer
        , CommandInfo -> LastKeyPositionSpec
    lastKeyPosition :: LastKeyPositionSpec
        , CommandInfo -> Integer
    stepCount :: Integer
        } deriving (Int -> CommandInfo -> ShowS
    [CommandInfo] -> ShowS
    CommandInfo -> [Char]
    (Int -> CommandInfo -> ShowS)
    -> (CommandInfo -> [Char])
    -> ([CommandInfo] -> ShowS)
    -> Show CommandInfo
    forall a.
    (Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> CommandInfo -> ShowS
    showsPrec :: Int -> CommandInfo -> ShowS
    $cshow :: CommandInfo -> [Char]
    show :: CommandInfo -> [Char]
    $cshowList :: [CommandInfo] -> ShowS
    showList :: [CommandInfo] -> ShowS
    Show)

    instance RedisResult CommandInfo where
        decode :: Reply -> Either Reply CommandInfo
    decode (MultiBulk (Just
            [ Bulk (Just ByteString
    commandName)
            , Integer Integer
    aritySpec
            , MultiBulk (Just [Reply]
    replyFlags)
            , Integer Integer
    firstKeyPos
            , Integer Integer
    lastKeyPos
            , Integer Integer
    replyStepCount])) = do
                parsedFlags <- (Reply -> Either Reply Flag) -> [Reply] -> Either Reply [Flag]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM Reply -> Either Reply Flag
    parseFlag [Reply]
    replyFlags
                lastKey <- parseLastKeyPos
                return $ CommandInfo
                    { name = commandName
                    , arity = parseArity aritySpec
                    , flags = parsedFlags
                    , firstKeyPosition = firstKeyPos
                    , lastKeyPosition = lastKey
                    , stepCount = replyStepCount
                    } where
            parseArity :: Integer -> AritySpec
    parseArity Integer
    int = case Integer
    int of
                Integer
    i | Integer
    i Integer -> Integer -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Integer
    0 -> Integer -> AritySpec
    Required Integer
    i
                Integer
    i -> Integer -> AritySpec
    MinimumRequired (Integer -> AritySpec) -> Integer -> AritySpec
    forall a b. (a -> b) -> a -> b
    $ Integer -> Integer
    forall a. Num a => a -> a
    abs Integer
    i
            parseFlag :: Reply -> Either Reply Flag
            parseFlag :: Reply -> Either Reply Flag
    parseFlag (SingleLine ByteString
    flag) = Flag -> Either Reply Flag
    forall a. a -> Either Reply a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (Flag -> Either Reply Flag) -> Flag -> Either Reply Flag
    forall a b. (a -> b) -> a -> b
    $ case ByteString
    flag of
                ByteString
    "write" -> Flag
    Write
                ByteString
    "readonly" -> Flag
    ReadOnly
                ByteString
    "denyoom" -> Flag
    DenyOOM
                ByteString
    "admin" -> Flag
    Admin
                ByteString
    "pubsub" -> Flag
    PubSub
                ByteString
    "noscript" -> Flag
    NoScript
                ByteString
    "random" -> Flag
    Random
                ByteString
    "sort_for_script" -> Flag
    SortForScript
                ByteString
    "loading" -> Flag
    Loading
                ByteString
    "stale" -> Flag
    Stale
                ByteString
    "skip_monitor" -> Flag
    SkipMonitor
                ByteString
    "asking" -> Flag
    Asking
                ByteString
    "fast" -> Flag
    Fast
                ByteString
    "movablekeys" -> Flag
    MovableKeys
                ByteString
    other -> ByteString -> Flag
    Other ByteString
    other
            parseFlag Reply
    bad = Reply -> Either Reply Flag
    forall a b. a -> Either a b
    Left Reply
    bad
            parseLastKeyPos :: Either Reply LastKeyPositionSpec
            parseLastKeyPos :: Either Reply LastKeyPositionSpec
    parseLastKeyPos = LastKeyPositionSpec -> Either Reply LastKeyPositionSpec
    forall a. a -> Either Reply a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (LastKeyPositionSpec -> Either Reply LastKeyPositionSpec)
    -> LastKeyPositionSpec -> Either Reply LastKeyPositionSpec
    forall a b. (a -> b) -> a -> b
    $ case Integer
    lastKeyPos of
                Integer
    i | Integer
    i Integer -> Integer -> Bool
    forall a. Ord a => a -> a -> Bool
    < Integer
    0 -> Integer -> LastKeyPositionSpec
    UnlimitedKeys (-Integer
    i Integer -> Integer -> Integer
    forall a. Num a => a -> a -> a
    - Integer
    1)
                Integer
    i -> Integer -> LastKeyPositionSpec
    LastKeyPosition Integer
    i
        -- since redis 6.0
        decode (MultiBulk (Just
            [ name :: Reply
    name@(Bulk (Just ByteString
    _))
            , arity :: Reply
    arity@(Integer Integer
    _)
            , flags :: Reply
    flags@(MultiBulk (Just [Reply]
    _))
            , firstPos :: Reply
    firstPos@(Integer Integer
    _)
            , lastPos :: Reply
    lastPos@(Integer Integer
    _)
            , step :: Reply
    step@(Integer Integer
    _)
            , MultiBulk Maybe [Reply]
    _  -- ACL categories
            ])) =
            Reply -> Either Reply CommandInfo
    forall a. RedisResult a => Reply -> Either Reply a
    decode (Maybe [Reply] -> Reply
    MultiBulk ([Reply] -> Maybe [Reply]
    forall a. a -> Maybe a
    Just [Reply
    name, Reply
    arity, Reply
    flags, Reply
    firstPos, Reply
    lastPos, Reply
    step]))

        decode Reply
    e = Reply -> Either Reply CommandInfo
    forall a b. a -> Either a b
    Left Reply
    e

    newInfoMap :: [CommandInfo] -> InfoMap
    newInfoMap :: [CommandInfo] -> InfoMap
    newInfoMap = HashMap [Char] CommandInfo -> InfoMap
    InfoMap (HashMap [Char] CommandInfo -> InfoMap)
    -> ([CommandInfo] -> HashMap [Char] CommandInfo)
    -> [CommandInfo]
    -> InfoMap
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [([Char], CommandInfo)] -> HashMap [Char] CommandInfo
    forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
    HM.fromList ([([Char], CommandInfo)] -> HashMap [Char] CommandInfo)
    -> ([CommandInfo] -> [([Char], CommandInfo)])
    -> [CommandInfo]
    -> HashMap [Char] CommandInfo
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (CommandInfo -> ([Char], CommandInfo))
    -> [CommandInfo] -> [([Char], CommandInfo)]
    forall a b. (a -> b) -> [a] -> [b]
    map (\CommandInfo
    c -> (ByteString -> [Char]
    Char8.unpack (ByteString -> [Char]) -> ByteString -> [Char]
    forall a b. (a -> b) -> a -> b
    $ CommandInfo -> ByteString
    name CommandInfo
    c, CommandInfo
    c))

    keysForRequest :: InfoMap -> [BS.ByteString] -> Maybe [BS.ByteString]
    keysForRequest :: InfoMap -> [ByteString] -> Maybe [ByteString]
    keysForRequest InfoMap
    _ [ByteString
    "DEBUG", ByteString
    "OBJECT", ByteString
    key] =
        -- `COMMAND` output for `DEBUG` would let us believe it doesn't have any
        -- keys, but the `DEBUG OBJECT` subcommand does.
        [ByteString] -> Maybe [ByteString]
    forall a. a -> Maybe a
    Just [ByteString
    key]
    keysForRequest InfoMap
    _ [ByteString
    "QUIT"] =
        -- The `QUIT` command is not listed in the `COMMAND` output.
        [ByteString] -> Maybe [ByteString]
    forall a. a -> Maybe a
    Just []
    keysForRequest (InfoMap HashMap [Char] CommandInfo
    infoMap) request :: [ByteString]
    request@(ByteString
    command:[ByteString]
    _) = do
        info <- [Char] -> HashMap [Char] CommandInfo -> Maybe CommandInfo
    forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
    HM.lookup ((Char -> Char) -> ShowS
    forall a b. (a -> b) -> [a] -> [b]
    map Char -> Char
    toLower ShowS -> ShowS
    forall a b. (a -> b) -> a -> b
    $ ByteString -> [Char]
    Char8.unpack ByteString
    command) HashMap [Char] CommandInfo
    infoMap
        keysForRequest' info request
    keysForRequest InfoMap
    _ [] = Maybe [ByteString]
    forall a. Maybe a
    Nothing

    keysForRequest' :: CommandInfo -> [BS.ByteString] -> Maybe [BS.ByteString]
    keysForRequest' :: CommandInfo -> [ByteString] -> Maybe [ByteString]
    keysForRequest' CommandInfo
    info [ByteString]
    request
        | CommandInfo -> Bool
    isMovable CommandInfo
    info =
            [ByteString] -> Maybe [ByteString]
    parseMovable [ByteString]
    request
        | CommandInfo -> Integer
    stepCount CommandInfo
    info Integer -> Integer -> Bool
    forall a. Eq a => a -> a -> Bool
    == Integer
    0 =
            [ByteString] -> Maybe [ByteString]
    forall a. a -> Maybe a
    Just []
        | Bool
    otherwise = do
            let possibleKeys :: [ByteString]
    possibleKeys = case CommandInfo -> LastKeyPositionSpec
    lastKeyPosition CommandInfo
    info of
                    LastKeyPosition Integer
    end -> Int -> [ByteString] -> [ByteString]
    forall a. Int -> [a] -> [a]
    take (Integer -> Int
    forall a. Enum a => a -> Int
    fromEnum (Integer -> Int) -> Integer -> Int
    forall a b. (a -> b) -> a -> b
    $ Integer
    1 Integer -> Integer -> Integer
    forall a. Num a => a -> a -> a
    + Integer
    end Integer -> Integer -> Integer
    forall a. Num a => a -> a -> a
    - CommandInfo -> Integer
    firstKeyPosition CommandInfo
    info) ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
    forall a b. (a -> b) -> a -> b
    $ Int -> [ByteString] -> [ByteString]
    forall a. Int -> [a] -> [a]
    drop (Integer -> Int
    forall a. Enum a => a -> Int
    fromEnum (Integer -> Int) -> Integer -> Int
    forall a b. (a -> b) -> a -> b
    $ CommandInfo -> Integer
    firstKeyPosition CommandInfo
    info) [ByteString]
    request
                    UnlimitedKeys Integer
    end ->
                        Int -> [ByteString] -> [ByteString]
    forall a. Int -> [a] -> [a]
    drop (Integer -> Int
    forall a. Enum a => a -> Int
    fromEnum (Integer -> Int) -> Integer -> Int
    forall a b. (a -> b) -> a -> b
    $ CommandInfo -> Integer
    firstKeyPosition CommandInfo
    info) ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
    forall a b. (a -> b) -> a -> b
    $
                           Int -> [ByteString] -> [ByteString]
    forall a. Int -> [a] -> [a]
    take ([ByteString] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [ByteString]
    request Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Integer -> Int
    forall a. Enum a => a -> Int
    fromEnum Integer
    end) [ByteString]
    request
            [ByteString] -> Maybe [ByteString]
    forall a. a -> Maybe a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ([ByteString] -> Maybe [ByteString])
    -> [ByteString] -> Maybe [ByteString]
    forall a b. (a -> b) -> a -> b
    $ Int -> [ByteString] -> [ByteString]
    forall a. Int -> [a] -> [a]
    takeEvery (Integer -> Int
    forall a. Enum a => a -> Int
    fromEnum (Integer -> Int) -> Integer -> Int
    forall a b. (a -> b) -> a -> b
    $ CommandInfo -> Integer
    stepCount CommandInfo
    info) [ByteString]
    possibleKeys

    isMovable :: CommandInfo -> Bool
    isMovable :: CommandInfo -> Bool
    isMovable CommandInfo{Integer
    [Flag]
    ByteString
    LastKeyPositionSpec
    AritySpec
    name :: CommandInfo -> ByteString
    arity :: CommandInfo -> AritySpec
    flags :: CommandInfo -> [Flag]
    firstKeyPosition :: CommandInfo -> Integer
    lastKeyPosition :: CommandInfo -> LastKeyPositionSpec
    stepCount :: CommandInfo -> Integer
    name :: ByteString
    arity :: AritySpec
    flags :: [Flag]
    firstKeyPosition :: Integer
    lastKeyPosition :: LastKeyPositionSpec
    stepCount :: Integer
    ..} = Flag
    MovableKeys Flag -> [Flag] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` [Flag]
    flags

    parseMovable :: [BS.ByteString] -> Maybe [BS.ByteString]
    parseMovable :: [ByteString] -> Maybe [ByteString]
    parseMovable (ByteString
    "SORT":ByteString
    key:[ByteString]
    _) = [ByteString] -> Maybe [ByteString]
    forall a. a -> Maybe a
    Just [ByteString
    key]
    parseMovable (ByteString
    "EVAL":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readNumKeys [ByteString]
    rest
    parseMovable (ByteString
    "EVALSHA":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readNumKeys [ByteString]
    rest
    parseMovable (ByteString
    "ZUNIONSTORE":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readNumKeys [ByteString]
    rest
    parseMovable (ByteString
    "ZINTERSTORE":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readNumKeys [ByteString]
    rest
    parseMovable (ByteString
    "XREAD":[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readXreadKeys [ByteString]
    rest
    parseMovable (ByteString
    "XREADGROUP":ByteString
    "GROUP":ByteString
    _:ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readXreadgroupKeys [ByteString]
    rest
    parseMovable [ByteString]
    _ = Maybe [ByteString]
    forall a. Maybe a
    Nothing

    readXreadKeys :: [BS.ByteString] -> Maybe [BS.ByteString]
    readXreadKeys :: [ByteString] -> Maybe [ByteString]
    readXreadKeys (ByteString
    "COUNT":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readXreadKeys [ByteString]
    rest
    readXreadKeys (ByteString
    "BLOCK":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readXreadKeys [ByteString]
    rest
    readXreadKeys (ByteString
    "STREAMS":[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    forall a. a -> Maybe a
    Just ([ByteString] -> Maybe [ByteString])
    -> [ByteString] -> Maybe [ByteString]
    forall a b. (a -> b) -> a -> b
    $ Int -> [ByteString] -> [ByteString]
    forall a. Int -> [a] -> [a]
    take ([ByteString] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [ByteString]
    rest Int -> Int -> Int
    forall a. Integral a => a -> a -> a
    `div` Int
    2) [ByteString]
    rest
    readXreadKeys [ByteString]
    _ = Maybe [ByteString]
    forall a. Maybe a
    Nothing

    readXreadgroupKeys :: [BS.ByteString] -> Maybe [BS.ByteString]
    readXreadgroupKeys :: [ByteString] -> Maybe [ByteString]
    readXreadgroupKeys (ByteString
    "COUNT":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readXreadKeys [ByteString]
    rest
    readXreadgroupKeys (ByteString
    "BLOCK":ByteString
    _:[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readXreadKeys [ByteString]
    rest
    readXreadgroupKeys (ByteString
    "NOACK":[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    readXreadKeys [ByteString]
    rest
    readXreadgroupKeys (ByteString
    "STREAMS":[ByteString]
    rest) = [ByteString] -> Maybe [ByteString]
    forall a. a -> Maybe a
    Just ([ByteString] -> Maybe [ByteString])
    -> [ByteString] -> Maybe [ByteString]
    forall a b. (a -> b) -> a -> b
    $ Int -> [ByteString] -> [ByteString]
    forall a. Int -> [a] -> [a]
    take ([ByteString] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [ByteString]
    rest Int -> Int -> Int
    forall a. Integral a => a -> a -> a
    `div` Int
    2) [ByteString]
    rest
    readXreadgroupKeys [ByteString]
    _ = Maybe [ByteString]
    forall a. Maybe a
    Nothing

    readNumKeys :: [BS.ByteString] -> Maybe [BS.ByteString]
    readNumKeys :: [ByteString] -> Maybe [ByteString]
    readNumKeys (ByteString
    rawNumKeys:[ByteString]
    rest) = do
        numKeys <- [Char] -> Maybe Int
    forall a. Read a => [Char] -> Maybe a
    readMaybe (ByteString -> [Char]
    Char8.unpack ByteString
    rawNumKeys)
        return $ take numKeys rest
    readNumKeys [ByteString]
    _ = Maybe [ByteString]
    forall a. Maybe a
    Nothing
    -- takeEvery 1 [1,2,3,4,5] ->[1,2,3,4,5]
    -- takeEvery 2 [1,2,3,4,5] ->[1,3,5]
    -- takeEvery 3 [1,2,3,4,5] ->[1,4]
    takeEvery :: Int -> [a] -> [a]
    takeEvery :: forall a. Int -> [a] -> [a]
    takeEvery Int
    _ [] = []
    takeEvery Int
    n (a
    x:[a]
    xs) = a
    x a -> [a] -> [a]
    forall a. a -> [a] -> [a]
    : Int -> [a] -> [a]
    forall a. Int -> [a] -> [a]
    takeEvery Int
    n (Int -> [a] -> [a]
    forall a. Int -> [a] -> [a]
    drop (Int
    nInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1) [a]
    xs)

    readMaybe :: Read a => String -> Maybe a
    readMaybe :: forall a. Read a => [Char] -> Maybe a
    readMaybe [Char]
    s = case ReadS a
    forall a. Read a => ReadS a
    reads [Char]
    s of
                      [(a
    val, [Char]
    "")] -> a -> Maybe a
    forall a. a -> Maybe a
    Just a
    val
                      [(a, [Char])]
    _           -> Maybe a
    forall a. Maybe a
    Nothing