module Redis::Commands::SortedSets

def _zsets_operation(cmd, *keys, weights: nil, aggregate: nil, with_scores: false)
  keys.flatten!(1)
  command = [cmd, keys.size].concat(keys)
  if weights
    command << "WEIGHTS"
    command.concat(weights)
  end
  command << "AGGREGATE" << aggregate if aggregate
  if with_scores
    command << "WITHSCORES"
    block = FloatifyPairs
  end
  send_command(command, &block)
end

def _zsets_operation_store(cmd, destination, keys, weights: nil, aggregate: nil)
  keys.flatten!(1)
  command = [cmd, destination, keys.size].concat(keys)
  if weights
    command << "WEIGHTS"
    command.concat(weights)
  end
  command << "AGGREGATE" << aggregate if aggregate
  send_command(command)
end

def bzpopmax(*args)
  _bpop(:bzpopmax, args) do |reply|
    reply.is_a?(Array) ? [reply[0], reply[1], Floatify.call(reply[2])] : reply
  end
end

def bzpopmin(*args)
  _bpop(:bzpopmin, args) do |reply|
    reply.is_a?(Array) ? [reply[0], reply[1], Floatify.call(reply[2])] : reply
  end
end

def zadd(key, *args, nx: nil, xx: nil, lt: nil, gt: nil, ch: nil, incr: nil)
  command = [:zadd, key]
  command << "NX" if nx
  command << "XX" if xx
  command << "LT" if lt
  command << "GT" if gt
  command << "CH" if ch
  command << "INCR" if incr
  if args.size == 1 && args[0].is_a?(Array)
    members_to_add = args[0]
    return 0 if members_to_add.empty?
    # Variadic: return float if INCR, integer if !INCR
    send_command(command + members_to_add, &(incr ? Floatify : nil))
  elsif args.size == 2
    # Single pair: return float if INCR, boolean if !INCR
    send_command(command + args, &(incr ? Floatify : Boolify))
  else
    raise ArgumentError, "wrong number of arguments"
  end
end

def zcard(key)
  send_command([:zcard, key])
end

def zcount(key, min, max)
  send_command([:zcount, key, min, max])
end

def zdiff(*keys, with_scores: false)
  _zsets_operation(:zdiff, *keys, with_scores: with_scores)
end

def zdiffstore(*args)
  _zsets_operation_store(:zdiffstore, *args)
end

def zincrby(key, increment, member)
  send_command([:zincrby, key, increment, member], &Floatify)
end

def zinter(*args)
  _zsets_operation(:zinter, *args)
end

def zinterstore(*args)
  _zsets_operation_store(:zinterstore, *args)
end

def zlexcount(key, min, max)
  send_command([:zlexcount, key, min, max])
end

def zmscore(key, *members)
  send_command([:zmscore, key, *members]) do |reply|
    reply.map(&Floatify)
  end
end

def zpopmax(key, count = nil)
  command = [:zpopmax, key]
  command << Integer(count) if count
  send_command(command) do |members|
    members = FloatifyPairs.call(members)
    count.to_i > 1 ? members : members.first
  end
end

def zpopmin(key, count = nil)
  command = [:zpopmin, key]
  command << Integer(count) if count
  send_command(command) do |members|
    members = FloatifyPairs.call(members)
    count.to_i > 1 ? members : members.first
  end
end

def zrandmember(key, count = nil, withscores: false, with_scores: withscores)
  if with_scores && count.nil?
    raise ArgumentError, "count argument must be specified"
  end
  args = [:zrandmember, key]
  args << Integer(count) if count
  if with_scores
    args << "WITHSCORES"
    block = FloatifyPairs
  end
  send_command(args, &block)
end

def zrange(key, start, stop, byscore: false, by_score: byscore, bylex: false, by_lex: bylex,
           rev: false, limit: nil, withscores: false, with_scores: withscores)
  if by_score && by_lex
    raise ArgumentError, "only one of :by_score or :by_lex can be specified"
  end
  args = [:zrange, key, start, stop]
  if by_score
    args << "BYSCORE"
  elsif by_lex
    args << "BYLEX"
  end
  args << "REV" if rev
  if limit
    args << "LIMIT"
    args.concat(limit.map { |l| Integer(l) })
  end
  if with_scores
    args << "WITHSCORES"
    block = FloatifyPairs
  end
  send_command(args, &block)
end

def zrangebylex(key, min, max, limit: nil)
  args = [:zrangebylex, key, min, max]
  if limit
    args << "LIMIT"
    args.concat(limit.map { |l| Integer(l) })
  end
  send_command(args)
end

def zrangebyscore(key, min, max, withscores: false, with_scores: withscores, limit: nil)
  args = [:zrangebyscore, key, min, max]
  if with_scores
    args << "WITHSCORES"
    block = FloatifyPairs
  end
  if limit
    args << "LIMIT"
    args.concat(limit.map { |l| Integer(l) })
  end
  send_command(args, &block)
end

def zrangestore(dest_key, src_key, start, stop, byscore: false, by_score: byscore,
                bylex: false, by_lex: bylex, rev: false, limit: nil)
  if by_score && by_lex
    raise ArgumentError, "only one of :by_score or :by_lex can be specified"
  end
  args = [:zrangestore, dest_key, src_key, start, stop]
  if by_score
    args << "BYSCORE"
  elsif by_lex
    args << "BYLEX"
  end
  args << "REV" if rev
  if limit
    args << "LIMIT"
    args.concat(limit.map { |l| Integer(l) })
  end
  send_command(args)
end

def zrank(key, member)
  send_command([:zrank, key, member])
end

def zrem(key, member)
  if member.is_a?(Array)
    members_to_remove = member
    return 0 if members_to_remove.empty?
  end
  send_command([:zrem, key, member]) do |reply|
    if member.is_a? Array
      # Variadic: return integer
      reply
    else
      # Single argument: return boolean
      Boolify.call(reply)
    end
  end
end

def zremrangebyrank(key, start, stop)
  send_command([:zremrangebyrank, key, start, stop])
end

def zremrangebyscore(key, min, max)
  send_command([:zremrangebyscore, key, min, max])
end

def zrevrange(key, start, stop, withscores: false, with_scores: withscores)
  args = [:zrevrange, key, Integer(start), Integer(stop)]
  if with_scores
    args << "WITHSCORES"
    block = FloatifyPairs
  end
  send_command(args, &block)
end

def zrevrangebylex(key, max, min, limit: nil)
  args = [:zrevrangebylex, key, max, min]
  if limit
    args << "LIMIT"
    args.concat(limit.map { |l| Integer(l) })
  end
  send_command(args)
end

def zrevrangebyscore(key, max, min, withscores: false, with_scores: withscores, limit: nil)
  args = [:zrevrangebyscore, key, max, min]
  if with_scores
    args << "WITHSCORES"
    block = FloatifyPairs
  end
  if limit
    args << "LIMIT"
    args.concat(limit.map { |l| Integer(l) })
  end
  send_command(args, &block)
end

def zrevrank(key, member)
  send_command([:zrevrank, key, member])
end

def zscan(key, cursor, **options)
  _scan(:zscan, cursor, [key], **options) do |reply|
    [reply[0], FloatifyPairs.call(reply[1])]
  end
end

def zscan_each(key, **options, &block)
  return to_enum(:zscan_each, key, **options) unless block_given?
  cursor = 0
  loop do
    cursor, values = zscan(key, cursor, **options)
    values.each(&block)
    break if cursor == "0"
  end
end

def zscore(key, member)
  send_command([:zscore, key, member], &Floatify)
end

def zunion(*args)
  _zsets_operation(:zunion, *args)
end

def zunionstore(*args)
  _zsets_operation_store(:zunionstore, *args)
end