diff --git a/app/components/ruby_ui/base.rb b/app/components/ruby_ui/base.rb
index a60b8bb4..ae8bf8ae 100644
--- a/app/components/ruby_ui/base.rb
+++ b/app/components/ruby_ui/base.rb
@@ -18,5 +18,11 @@ def initialize(**user_attrs)
     def default_attrs
       {}
     end
+
+    if Rails.env.development?
+      def before_template
+        comment { "Before #{self.class.name}" }
+      end
+    end
   end
 end
diff --git a/app/components/ruby_ui/theme_toggle/set_dark_mode.rb b/app/components/ruby_ui/theme_toggle/set_dark_mode.rb
new file mode 100644
index 00000000..65019f06
--- /dev/null
+++ b/app/components/ruby_ui/theme_toggle/set_dark_mode.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module RubyUI
+  class SetDarkMode < ThemeToggle
+    def view_template(&)
+      div(**attrs, &)
+    end
+
+    def default_attrs
+      {
+        class: "hidden dark:inline-block",
+        data: {controller: "ruby-ui--theme-toggle", action: "click->ruby-ui--theme-toggle#setLightTheme"}
+      }
+    end
+  end
+end
diff --git a/app/components/ruby_ui/theme_toggle/set_light_mode.rb b/app/components/ruby_ui/theme_toggle/set_light_mode.rb
new file mode 100644
index 00000000..205b0224
--- /dev/null
+++ b/app/components/ruby_ui/theme_toggle/set_light_mode.rb
@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+module RubyUI
+  class SetLightMode < ThemeToggle
+    def view_template(&)
+      div(**attrs, &)
+    end
+
+    def default_attrs
+      {
+        class: "dark:hidden",
+        data: {controller: "ruby-ui--theme-toggle", action: "click->ruby-ui--theme-toggle#setDarkTheme"}
+      }
+    end
+  end
+end
diff --git a/app/components/ruby_ui/theme_toggle/theme_toggle.rb b/app/components/ruby_ui/theme_toggle/theme_toggle.rb
index a667c1aa..000f8054 100644
--- a/app/components/ruby_ui/theme_toggle/theme_toggle.rb
+++ b/app/components/ruby_ui/theme_toggle/theme_toggle.rb
@@ -5,37 +5,5 @@ class ThemeToggle < Base
     def view_template(&)
       div(**attrs, &)
     end
-
-    def light_mode(**user_attrs, &)
-      light_attrs = mix(default_light_attrs, user_attrs)
-      div(**light_attrs, &)
-    end
-
-    def dark_mode(**user_attrs, &)
-      dark_attrs = mix(default_dark_attrs, user_attrs)
-      div(**dark_attrs, &)
-    end
-
-    private
-
-    def default_attrs
-      {
-        data: {controller: "ruby-ui--theme-toggle"}
-      }
-    end
-
-    def default_light_attrs
-      {
-        class: "dark:hidden",
-        data: {action: "click->ruby-ui--theme-toggle#setDarkTheme"}
-      }
-    end
-
-    def default_dark_attrs
-      {
-        class: "hidden dark:inline-block",
-        data: {action: "click->ruby-ui--theme-toggle#setLightTheme"}
-      }
-    end
   end
 end
diff --git a/app/components/shared/navbar.rb b/app/components/shared/navbar.rb
index ca08e5bd..aa8f34da 100644
--- a/app/components/shared/navbar.rb
+++ b/app/components/shared/navbar.rb
@@ -26,8 +26,8 @@ def view_template
       end
 
       def dark_mode_toggle
-        RubyUI.ThemeToggle do |toggle|
-          toggle.light_mode do
+        ThemeToggle do
+          SetLightMode do
             Button(variant: :ghost, icon: true) do
               svg(
                 xmlns: "http://www.w3.org/2000/svg",
@@ -42,8 +42,7 @@ def dark_mode_toggle
               end
             end
           end
-
-          toggle.dark_mode do
+          SetDarkMode do
             Button(variant: :ghost, icon: true) do
               svg(
                 xmlns: "http://www.w3.org/2000/svg",
diff --git a/app/views/docs/getting_started/dark_mode.rb b/app/views/docs/getting_started/dark_mode.rb
index 2a2c5424..43eb84f3 100644
--- a/app/views/docs/getting_started/dark_mode.rb
+++ b/app/views/docs/getting_started/dark_mode.rb
@@ -49,7 +49,7 @@ def view_template
           render Docs::VisualCodeExample.new(title: "Toggle component", context: self) do
             <<~RUBY
               ThemeToggle do |toggle|
-                toggle.light_mode do
+                SetLightMode do
                   Button(variant: :outline, icon: true) do
                     svg(
                       xmlns: "http://www.w3.org/2000/svg",
@@ -65,7 +65,7 @@ def view_template
                   end
                 end
 
-                toggle.dark_mode do
+                SetDarkMode do
                   Button(variant: :outline, icon: true) do
                     svg(
                       xmlns: "http://www.w3.org/2000/svg",
diff --git a/app/views/docs/theme_toggle.rb b/app/views/docs/theme_toggle.rb
index 7c26615e..1740a924 100644
--- a/app/views/docs/theme_toggle.rb
+++ b/app/views/docs/theme_toggle.rb
@@ -12,7 +12,7 @@ def view_template
       render Docs::VisualCodeExample.new(title: "With icon", context: self) do
         <<~RUBY
           ThemeToggle do |toggle|
-            toggle.light_mode do
+            SetLightMode do
               Button(variant: :ghost, icon: true) do
                 svg(
                   xmlns: "http://www.w3.org/2000/svg",
@@ -28,7 +28,7 @@ def view_template
               end
             end
 
-            toggle.dark_mode do
+            SetDarkMode do
               Button(variant: :ghost, icon: true) do
                 svg(
                   xmlns: "http://www.w3.org/2000/svg",
@@ -52,11 +52,11 @@ def view_template
       render Docs::VisualCodeExample.new(title: "With text", context: self) do
         <<~RUBY
           ThemeToggle do |toggle|
-            toggle.light_mode do
+            SetLightMode do
               Button(variant: :primary) { "Light" }
             end
 
-            toggle.dark_mode do
+            SetDarkMode do
               Button(variant: :primary) { "Dark" }
             end
           end