diff --git a/lib/flushbar.dart b/lib/flushbar.dart
index 43c38ab..793adb1 100644
--- a/lib/flushbar.dart
+++ b/lib/flushbar.dart
@@ -119,7 +119,7 @@ class Flushbar<T> extends StatefulWidget {
   final bool shouldIconPulse;
 
   /// Use if you need an action from the user. [TextButton] is recommended here
-  final Widget? mainButton;
+  final Widget Function(Flushbar<T>)? mainButton;
 
   /// A callback that registers the user's click anywhere. An alternative to [mainButton]
   final OnTap? onTap;
@@ -235,8 +235,7 @@ class Flushbar<T> extends StatefulWidget {
       flushbar: this,
     ) as route.FlushbarRoute<T?>;
 
-    return await Navigator.of(context, rootNavigator: false)
-        .push(flushbarRoute as Route<T>);
+    return await Navigator.of(context, rootNavigator: false).push(flushbarRoute as Route<T>);
   }
 
   /// Dismisses the flushbar causing is to return a future containing [result].
@@ -294,8 +293,7 @@ class Flushbar<T> extends StatefulWidget {
   State createState() => _FlushbarState<T?>();
 }
 
-class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
-    with TickerProviderStateMixin {
+class _FlushbarState<K extends Object?> extends State<Flushbar<K>> with TickerProviderStateMixin {
   final Duration _pulseAnimationDuration = const Duration(seconds: 1);
   final Widget _emptyWidget = const SizedBox();
   final double _initialOpacity = 1.0;
@@ -320,10 +318,7 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
     _backgroundBoxKey = GlobalKey();
     _boxHeightCompleter = Completer<Size>();
 
-    assert(
-        widget.userInputForm != null ||
-            ((widget.message != null && widget.message!.isNotEmpty) ||
-                widget.messageText != null),
+    assert(widget.userInputForm != null || ((widget.message != null && widget.message!.isNotEmpty) || widget.messageText != null),
         'A message is mandatory if you are not using userInputForm. Set either a message or messageText');
 
     _isTitlePresent = (widget.title != null || widget.titleText != null);
@@ -365,16 +360,13 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
   }
 
   void _configureProgressIndicatorAnimation() {
-    if (widget.showProgressIndicator &&
-        widget.progressIndicatorController != null) {
-      _progressAnimation = CurvedAnimation(
-          curve: Curves.linear, parent: widget.progressIndicatorController!);
+    if (widget.showProgressIndicator && widget.progressIndicatorController != null) {
+      _progressAnimation = CurvedAnimation(curve: Curves.linear, parent: widget.progressIndicatorController!);
     }
   }
 
   void _configurePulseAnimation() {
-    _fadeController =
-        AnimationController(vsync: this, duration: _pulseAnimationDuration);
+    _fadeController = AnimationController(vsync: this, duration: _pulseAnimationDuration);
     _fadeAnimation = Tween(begin: _initialOpacity, end: _finalOpacity).animate(
       CurvedAnimation(
         parent: _fadeController!,
@@ -400,9 +392,7 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
     return Align(
       heightFactor: 1.0,
       child: Material(
-        color: widget.flushbarStyle == FlushbarStyle.FLOATING
-            ? Colors.transparent
-            : widget.backgroundColor,
+        color: widget.flushbarStyle == FlushbarStyle.FLOATING ? Colors.transparent : widget.backgroundColor,
         child: GestureDetector(
           onTap: () => widget.onTap?.call(widget),
           child: _getFlushbar(),
@@ -433,8 +423,7 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
               return ClipRRect(
                 borderRadius: widget.borderRadius ?? BorderRadius.zero,
                 child: BackdropFilter(
-                  filter: ImageFilter.blur(
-                      sigmaX: widget.barBlur, sigmaY: widget.barBlur),
+                  filter: ImageFilter.blur(sigmaX: widget.barBlur, sigmaY: widget.barBlur),
                   child: Container(
                     height: snapshot.data!.height,
                     width: snapshot.data!.width,
@@ -457,21 +446,16 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
   Widget _generateInputFlushbar() {
     return Container(
       key: _backgroundBoxKey,
-      constraints: widget.maxWidth != null
-          ? BoxConstraints(maxWidth: widget.maxWidth!)
-          : null,
+      constraints: widget.maxWidth != null ? BoxConstraints(maxWidth: widget.maxWidth!) : null,
       decoration: BoxDecoration(
         color: widget.backgroundColor,
         gradient: widget.backgroundGradient,
         boxShadow: widget.boxShadows,
         borderRadius: widget.borderRadius,
-        border: widget.borderColor != null
-            ? Border.all(color: widget.borderColor!, width: widget.borderWidth)
-            : null,
+        border: widget.borderColor != null ? Border.all(color: widget.borderColor!, width: widget.borderWidth) : null,
       ),
       child: Padding(
-        padding: const EdgeInsets.only(
-            left: 8.0, right: 8.0, bottom: 8.0, top: 16.0),
+        padding: const EdgeInsets.only(left: 8.0, right: 8.0, bottom: 8.0, top: 16.0),
         child: FocusScope(
           node: _focusNode,
           autofocus: true,
@@ -484,17 +468,13 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
   Widget _generateFlushbar() {
     return Container(
       key: _backgroundBoxKey,
-      constraints: widget.maxWidth != null
-          ? BoxConstraints(maxWidth: widget.maxWidth!)
-          : null,
+      constraints: widget.maxWidth != null ? BoxConstraints(maxWidth: widget.maxWidth!) : null,
       decoration: BoxDecoration(
         color: widget.backgroundColor,
         gradient: widget.backgroundGradient,
         boxShadow: widget.boxShadows,
         borderRadius: widget.borderRadius,
-        border: widget.borderColor != null
-            ? Border.all(color: widget.borderColor!, width: widget.borderWidth)
-            : null,
+        border: widget.borderColor != null ? Border.all(color: widget.borderColor!, width: widget.borderWidth) : null,
       ),
       child: Column(
         mainAxisSize: MainAxisSize.min,
@@ -707,12 +687,8 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
                 borderRadius: widget.borderRadius == null
                     ? null
                     : widget.textDirection == TextDirection.ltr
-                        ? BorderRadius.only(
-                            topLeft: widget.borderRadius!.topLeft,
-                            bottomLeft: widget.borderRadius!.bottomLeft)
-                        : BorderRadius.only(
-                            topRight: widget.borderRadius!.topRight,
-                            bottomRight: widget.borderRadius!.bottomRight),
+                        ? BorderRadius.only(topLeft: widget.borderRadius!.topLeft, bottomLeft: widget.borderRadius!.bottomLeft)
+                        : BorderRadius.only(topRight: widget.borderRadius!.topRight, bottomRight: widget.borderRadius!.bottomRight),
                 color: widget.leftBarIndicatorColor,
               ),
             );
@@ -743,25 +719,20 @@ class _FlushbarState<K extends Object?> extends State<Flushbar<K>>
     return widget.titleText ??
         Text(
           widget.title ?? '',
-          style: TextStyle(
-              fontSize: widget.titleSize ?? 16.0,
-              color: widget.titleColor ?? Colors.white,
-              fontWeight: FontWeight.bold),
+          style: TextStyle(fontSize: widget.titleSize ?? 16.0, color: widget.titleColor ?? Colors.white, fontWeight: FontWeight.bold),
         );
   }
 
   Text _getDefaultNotificationText() {
     return Text(
       widget.message ?? '',
-      style: TextStyle(
-          fontSize: widget.messageSize ?? 14.0,
-          color: widget.messageColor ?? Colors.white),
+      style: TextStyle(fontSize: widget.messageSize ?? 14.0, color: widget.messageColor ?? Colors.white),
     );
   }
 
   Widget? _getMainActionButton() {
     if (widget.mainButton != null) {
-      return widget.mainButton;
+      return widget.mainButton!(widget);
     } else {
       return null;
     }
diff --git a/lib/flushbar_helper.dart b/lib/flushbar_helper.dart
index 3b478ad..0be9c7b 100644
--- a/lib/flushbar_helper.dart
+++ b/lib/flushbar_helper.dart
@@ -3,10 +3,7 @@ import 'package:flutter/material.dart';
 
 class FlushbarHelper {
   /// Get a success notification flushbar.
-  static Flushbar createSuccess(
-      {required String message,
-      String? title,
-      Duration duration = const Duration(seconds: 3)}) {
+  static Flushbar createSuccess({required String message, String? title, Duration duration = const Duration(seconds: 3)}) {
     return Flushbar(
       title: title,
       message: message,
@@ -20,10 +17,7 @@ class FlushbarHelper {
   }
 
   /// Get an information notification flushbar
-  static Flushbar createInformation(
-      {required String message,
-      String? title,
-      Duration duration = const Duration(seconds: 3)}) {
+  static Flushbar createInformation({required String message, String? title, Duration duration = const Duration(seconds: 3)}) {
     return Flushbar(
       title: title,
       message: message,
@@ -38,10 +32,7 @@ class FlushbarHelper {
   }
 
   /// Get a error notification flushbar
-  static Flushbar createError(
-      {required String message,
-      String? title,
-      Duration duration = const Duration(seconds: 3)}) {
+  static Flushbar createError({required String message, String? title, Duration duration = const Duration(seconds: 3)}) {
     return Flushbar(
       title: title,
       message: message,
@@ -58,7 +49,7 @@ class FlushbarHelper {
   /// Get a flushbar that can receive a user action through a button.
   static Flushbar createAction(
       {required String message,
-      required Widget button,
+      required Widget Function(Flushbar<dynamic>) button,
       String? title,
       Duration duration = const Duration(seconds: 3)}) {
     return Flushbar(