From 6e4ea7675d76ddebe136f302f6328dd2e905f7f4 Mon Sep 17 00:00:00 2001
From: Ashley Wulber <48420062+wash2@users.noreply.github.com>
Date: Fri, 21 Mar 2025 20:25:42 -0400
Subject: [PATCH] fix(wallpaper): color picker functionality

---
 .../src/pages/desktop/wallpaper/mod.rs        | 58 +++++++++++++------
 page/src/section.rs                           | 10 ++--
 2 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/cosmic-settings/src/pages/desktop/wallpaper/mod.rs b/cosmic-settings/src/pages/desktop/wallpaper/mod.rs
index f752512..944ec22 100644
--- a/cosmic-settings/src/pages/desktop/wallpaper/mod.rs
+++ b/cosmic-settings/src/pages/desktop/wallpaper/mod.rs
@@ -23,12 +23,12 @@ use cosmic::widget::{
     settings, tab_bar, text, toggler,
 };
 use cosmic::{
-    iced::{window, Alignment, Color, Length},
-    surface,
+    Apply, Element, Task,
+    widget::{ColorPickerModel, color_picker::ColorPickerUpdate, icon},
 };
 use cosmic::{
-    widget::{color_picker::ColorPickerUpdate, icon, ColorPickerModel},
-    Apply, Element, Task,
+    iced::{Alignment, Color, Length, window},
+    surface,
 };
 use cosmic_bg_config::Source;
 use cosmic_settings_page::Section;
@@ -731,30 +731,40 @@ impl Page {
             }
 
             Message::ColorAdd(message) => {
-                if let ColorPickerUpdate::ActionFinished = message {
-                    let _res = self
-                        .color_model
-                        .update::<crate::app::Message>(ColorPickerUpdate::AppliedColor);
+                match message {
+                    ColorPickerUpdate::ActionFinished => {
+                        if let Some(color) = self.color_model.get_applied_color() {
+                            let ret = self
+                                .color_model
+                                .update::<crate::app::Message>(ColorPickerUpdate::ActionFinished);
+                            let color = wallpaper::Color::Single([color.r, color.g, color.b]);
 
-                    if let Some(color) = self.color_model.get_applied_color() {
-                        let color = wallpaper::Color::Single([color.r, color.g, color.b]);
+                            if let Choice::Color(c) = self.selection.active.clone() {
+                                if let Err(why) = self.config.remove_custom_color(&c) {
+                                    tracing::error!(?why, "could not set custom color");
+                                }
+                                self.selection.remove_custom_color(&c);
+                            }
+                            if let Err(why) = self.config.add_custom_color(color.clone()) {
+                                tracing::error!(?why, "could not set custom color");
+                            }
 
-                        if let Err(why) = self.config.add_custom_color(color.clone()) {
-                            tracing::error!(?why, "could not set custom color");
+                            self.cached_display_handle = None;
+                            self.selection.replace_active_custom(color.clone());
+                            self.config_apply();
+
+                            return ret;
                         }
-
-                        self.selection.add_custom_color(color.clone());
-                        self.selection.active = Choice::Color(color);
-                        self.cached_display_handle = None;
-                        self.context_view = None;
+                    }
+                    m => {
+                        return self.color_model.update::<crate::app::Message>(m);
                     }
                 };
-
-                return self.color_model.update::<crate::app::Message>(message);
             }
 
             Message::ColorAddContext => {
                 self.context_view = Some(ContextView::AddColor);
+                self.selection.active = Choice::Color(wallpaper::Color::Single([0., 0., 0.]));
                 return cosmic::task::message(crate::app::Message::OpenContextDrawer(
                     self.entity,
                     fl!("color-picker").into(),
@@ -1112,6 +1122,16 @@ impl Context {
         }
     }
 
+    fn replace_active_custom(&mut self, color: wallpaper::Color) {
+        if let Choice::Color(active) = &self.active {
+            self.custom_colors.retain(|c| c != active);
+        }
+        if !self.custom_colors.contains(&color) {
+            self.custom_colors.push(color.clone());
+            self.active = Choice::Color(color);
+        }
+    }
+
     fn add_custom_image(&mut self, path: PathBuf, display: Image, selection: ImageHandle) {
         let key = self.paths.insert(path);
         self.is_custom.insert(key, ());
diff --git a/page/src/section.rs b/page/src/section.rs
index 583f828..dd4494e 100644
--- a/page/src/section.rs
+++ b/page/src/section.rs
@@ -95,11 +95,11 @@ impl<Message: 'static> Section<Message> {
     pub fn view<Model: Page<Message>>(
         mut self,
         func: impl for<'a> Fn(
-                &'a Binder<Message>,
-                &'a Model,
-                &'a Section<Message>,
-            ) -> cosmic::Element<'a, Message>
-            + 'static,
+            &'a Binder<Message>,
+            &'a Model,
+            &'a Section<Message>,
+        ) -> cosmic::Element<'a, Message>
+        + 'static,
     ) -> Self {
         self.view_fn = Box::new(move |binder, model: &dyn Page<Message>, section| {
             let model = model.downcast_ref::<Model>().unwrap_or_else(|| {