module Test.ZSetCommands where

import Test.HUnit
import Control.Monad.Reader
import Data.Maybe
import Data.List
import Data.Map (Map(..), fromList, toList, unionWith, intersectionWith)
import qualified Data.Map as M
import Database.Redis.Redis
import Test.Setup
import Test.Utils

tests = TestList [TestLabel "zrange, zrevrange and scard" test_zrange_zrevrange_zcard,
                  TestLabel "zadd and zrem" test_zadd_zrem,
                  TestLabel "zscore" test_zscore,
                  TestLabel "zincrBy" test_zincrBy,
                  TestLabel "zrangebyscore" test_zrangebyscore,
                  TestLabel "zrevrangebyscore" test_zrevrangebyscore,
                  TestLabel "zcount" test_zcount,
                  TestLabel "zremrangebyscore" test_zremrangebyscore,
                  TestLabel "zrank and zrevrank" test_zrank_zrevrank,
                  TestLabel "zremrangebyrank" test_zremrangebyrank,
                  TestLabel "zunionStore" test_zunionStore,
                  TestLabel "zinterStore" test_zinterStore]

asZSet :: Reply String -> IO [(String, Double)]
asZSet r = fromRMultiBulk r >>= return . build . map fromJust . fromJust
    where build (a:b:z) = (a, read b) : build z
          build (a:[]) = error "unpaired element"
          build [] = []

zsort :: [(String, Double)] -> [(String, Double)]
zsort = sortBy (\(_, a) (_, b) -> compare a b)

test_zrange_zrevrange_zcard = TestCase $ testRedis $
    let expected = zsort $ zip ["1", "2", "3", "4", "5"] (reverse [1.0, 2.0, 3.0, 4.0, 5.0])
    in do r <- ask
          addZSet
          liftIO $ do zrange r "zset" takeAll True >>= asZSet >>= assertEqual "" expected
                      zrevrange r "zset" takeAll True >>= asZSet >>= assertEqual "" (reverse expected)
                      zcard r "zset" >>= fromRInt >>= assertEqual "" (length expected)
                      zcard r "no-such-key" >>= fromRInt >>= assertEqual "" 0
                      zrange r "no-such-key" takeAll True >>= asZSet >>= assertEqual "" []
                      zrevrange r "no-such-key" takeAll True >>= asZSet >>= assertEqual "" []
                      zrange r "zset" (1,2) True >>= asZSet >>= assertEqual "" (take 2 $ drop 1 $ expected)
                      zrevrange r "zset" (1,2) True >>= asZSet >>= assertEqual "" (take 2 $ drop 1 $ reverse expected)
                      zrange r "zset" (1,2) False >>= fromRMultiBulk >>= return . map fromJust . fromJust >>= assertEqual "" (map fst $ take 2 $ drop 1 $ expected)
                      zrevrange r "zset" (1,2) False >>= fromRMultiBulk >>= return . map fromJust . fromJust >>= assertEqual "" (map fst $ take 2 $ drop 1 $ reverse expected)

test_zadd_zrem = TestCase $ testRedis $
    do r <- ask
       addZSet
       liftIO $ do zadd r "zset" 0.5 "6" >>= fromRInt >>= assertEqual "" 1
                   zadd r "zset" 0.6 "1" >>= fromRInt >>= assertEqual "element is allready in set" 0
                   z <- zrange r "zset" takeAll True >>= asZSet
                   assertEqual "" ("6", 0.5) $ head z
                   assertEqual "" ("1", 0.6) $ head $ tail z
                   zrem r "zset" "6" >>= fromRInt >>= assertEqual "" 1
                   zrem r "zset" "6" >>= fromRInt >>= assertEqual "" 0
                   z <- zrange r "zset" takeAll True >>= asZSet
                   assertEqual "" ("1", 0.6) $ head z

test_zscore = TestCase $ testRedis $
    let lookupBy _ _ [] = Nothing
        lookupBy f a (l:ls) = if f l == a
                              then Just l
                              else lookupBy f a ls
    in do r <- ask
          addZSet
          liftIO $ do z <- zrange r "zset" takeAll True >>= asZSet
                      let expected = snd `fmap` lookupBy fst "1" z
                      score <- zscore r "zset" "1" >>= fromRBulk :: IO (Maybe Double)
                      assertEqual "" expected score

test_zincrBy = TestCase $ testRedis $
    do r <- ask
       addZSet
       liftIO $ do zincrBy r "zset" 0.5 "5" >>= fromRBulk >>= assertEqual "" (Just (1.5 :: Double))
                   score <- zscore r "zset" "5"  >>= fromRBulk
                   assertEqual "" (Just (1.5 :: Double)) score
                   zincrBy r "zset" (-0.5) "5" >>= fromRBulk >>= assertEqual "" (Just (1.0 :: Double))
                   score <- zscore r "zset" "5"  >>= fromRBulk
                   assertEqual "" (Just (1.0 :: Double)) score

test_zrangebyscore = TestCase $ testRedis $
    do r <- ask
       addZSet
       liftIO $ do z <- zrange r "zset" takeAll True >>= asZSet
                   z' <- zrangebyscore r "zset" [1, 4] (Just (1, 2)) True >>= asZSet
                   let expected = take 2 $ drop 1 $ filter ((\x -> x >= 1 && x <= 4) . snd) z
                   assertEqual "" expected z'
                   z' <- zrangebyscore r "zset" (1 :: Double, 4 :: Double) (Just (1, 2)) True >>= asZSet
                   let expected = take 2 $ drop 1 $ filter ((\x -> x > 1 && x < 4) . snd) z
                   assertEqual "" expected z'

test_zrevrangebyscore = TestCase $ testRedis $
    do r <- ask
       addZSet
       liftIO $ do z <- zrange r "zset" takeAll True >>= asZSet
                   z' <- zrevrangebyscore r "zset" [4, 1] (Just (1, 2)) True >>= asZSet
                   let expected = take 2 $ drop 1 $ reverse $ filter ((\x -> x >= 1 && x <= 4) . snd) z
                   assertEqual "[1, 4]" expected z'
                   z' <- zrevrangebyscore r "zset" (4 :: Double, 1 :: Double) (Just (1, 2)) True >>= asZSet
                   let expected = take 2 $ drop 1 $ reverse $ filter ((\x -> x > 1 && x < 4) . snd) z
                   assertEqual "(1, 4)" expected z'

test_zcount = TestCase $ testRedis $
    do r <- ask
       addZSet
       liftIO $ do z <- zrange r "zset" takeAll True >>= asZSet
                   z' <- zcount r "zset" [1, 4] >>= fromRInt
                   let expected = filter ((\x -> x >= 1 && x <= 4) . snd) z
                   assertEqual "" (length expected) z'
                   z' <- zcount r "zset" (1 :: Double, 4 :: Double) >>= fromRInt
                   let expected = filter ((\x -> x > 1 && x < 4) . snd) z
                   assertEqual "" (length expected) z'

test_zremrangebyscore = TestCase $ testRedis $
    do r <- ask
       addZSet
       liftIO $ do z <- zrange r "zset" takeAll True >>= asZSet
                   let removed = filter ((\x -> x >= 1 && x <= 4) . snd) z
                       left = filter (not . (\x -> x >= 1 && x <= 4) . snd) z
                   zremrangebyscore r "zset" (1, 4) >>= fromRInt >>= assertEqual "" (length removed)
                   z' <- zrange r "zset" takeAll True >>= asZSet
                   assertEqual "" left z'

test_zrank_zrevrank = TestCase $ testRedis $
    let rank e = findIndex (== e)
    in do r <- ask
          addZSet
          liftIO $ do z <- zrange r "zset" takeAll False >>= fromRMultiBulk >>= return . fromJust
                      c <- zrank r "zset" "3" >>= fromRInt
                      assertEqual "" (fromJust $ rank (Just "3") z) c
                      c <- zrevrank r "zset" "3" >>= fromRInt
                      assertEqual "" (fromJust $ rank (Just "3") $ reverse z) c

                      {- The next one doesn't works because of zrank returning
                         "RBulk Nil" if requested element is not a
                         member of set
                         zrank r "zset" "6" >>= fromRInt
                       -}

test_zremrangebyrank = TestCase $ testRedis $
    do r <- ask
       addZSet
       liftIO $ do z <- zrange r "zset" takeAll False >>= fromRMultiBulk >>= return . fromJust :: IO [Maybe String]
                   let expected = (take 1 z) ++ (drop 4 $ z) -- cut elements from 1 to 3
                   zremrangebyrank r "zset" (1, 3) >>= fromRInt >>= assertEqual "" (length z - length expected)
                   z' <- zrange r "zset" takeAll False >>= fromRMultiBulk >>= return . fromJust :: IO [Maybe String]
                   assertEqual "" expected z'

test_zunionStore = TestCase $ testRedis $
    let addZSet2 key m = ask >>= \r ->  liftIO $ mapM_ (uncurry (flip $ zadd r key)) $ toList m
        zmap2 = fromList $ zip ["1", "2", "6"] [1.0, 2.0, 3.0] :: Map String Double
    in do r <- ask
          addZSet
          addZSet2 "zset2" zmap2
          liftIO $ do z <- zrange r "zset" takeAll True >>= asZSet
                      let zmap = fromList z :: Map String Double
                          expected = zsort $ toList $ unionWith (+) zmap zmap2
                      res <- zunionStore r "zset3" ["zset", "zset2"] [1, 1] SUM >>= fromRInt
                      assertEqual "" (length expected) res
                      z' <- zrange r "zset3" takeAll True >>= asZSet
                      assertEqual "" expected z'

                      let zmap = fromList z :: Map String Double
                          -- in Redis weights is applied to maps before aggregation
                          expected = zsort $ toList $ unionWith (+) zmap $ M.map (* 0.5) zmap2
                      res <- zunionStore r "zset3" ["zset", "zset2"] [1, 0.5] SUM >>= fromRInt
                      assertEqual "" (length expected) res
                      z' <- zrange r "zset3" takeAll True >>= asZSet
                      assertEqual "" expected z'

test_zinterStore = TestCase $ testRedis $
    let addZSet2 key m = ask >>= \r -> liftIO $ mapM_ (uncurry (flip $ zadd r key)) $ toList m
        zmap2 = fromList $ zip ["1", "2", "6"] [1.0, 2.0, 3.0] :: Map String Double
    in do r <- ask
          addZSet
          addZSet2 "zset2" zmap2
          liftIO $ do z <- zrange r "zset" takeAll True >>= asZSet
                      let zmap = fromList z :: Map String Double
                          expected = zsort $ toList $ intersectionWith (+) zmap zmap2
                      res <- zinterStore r "zset3" ["zset", "zset2"] [1, 1] SUM >>= fromRInt
                      assertEqual "" (length expected) res
                      z' <- zrange r "zset3" takeAll True >>= asZSet
                      assertEqual "" expected z'

                      let zmap = fromList z :: Map String Double
                          -- in Redis weights is applied to maps before aggregation
                          expected = zsort $ toList $ intersectionWith (+) zmap $ M.map (* 0.5) zmap2
                      res <- zinterStore r "zset3" ["zset", "zset2"] [1, 0.5] SUM >>= fromRInt
                      assertEqual "" (length expected) res
                      z' <- zrange r "zset3" takeAll True >>= asZSet
                      assertEqual "" expected z'