From 66491a4e73d157649552cfb9f8d962a1ab194945 Mon Sep 17 00:00:00 2001
From: Teoh Han Hui <teohhanhui@gmail.com>
Date: Mon, 8 Jul 2024 03:39:46 +0800
Subject: [PATCH] Implement recoloring of ascii art using "custom" colors

---
 crates/hyfetch/src/models.rs        | 15 +++--
 crates/hyfetch/src/neofetch_util.rs | 25 +++++++-
 crates/hyfetch/src/utils.rs         | 94 +++++++++++++++++++++++++++++
 3 files changed, 124 insertions(+), 10 deletions(-)

diff --git a/crates/hyfetch/src/models.rs b/crates/hyfetch/src/models.rs
index 7751d2a5..10265421 100644
--- a/crates/hyfetch/src/models.rs
+++ b/crates/hyfetch/src/models.rs
@@ -14,7 +14,7 @@ pub struct Config {
     pub color_align: ColorAlignment,
     pub backend: Backend,
     #[serde(default)]
-    #[serde(with = "self::args_serde_with")]
+    #[serde(with = "self::args_serde")]
     pub args: Option<Vec<String>>,
     pub distro: Option<String>,
     pub pride_month_disable: bool,
@@ -38,16 +38,15 @@ impl Config {
     }
 }
 
-mod args_serde_with {
+mod args_serde {
     use std::fmt;
 
     use serde::de::{self, value, Deserialize, Deserializer, SeqAccess, Visitor};
     use serde::ser::Serializer;
 
-    pub(super) fn serialize<S>(
-        value: &Option<Vec<String>>,
-        serializer: S,
-    ) -> Result<S::Ok, S::Error>
+    type Value = Option<Vec<String>>;
+
+    pub(super) fn serialize<S>(value: &Value, serializer: S) -> Result<S::Ok, S::Error>
     where
         S: Serializer,
     {
@@ -57,7 +56,7 @@ mod args_serde_with {
         }
     }
 
-    pub(super) fn deserialize<'de, D>(deserializer: D) -> Result<Option<Vec<String>>, D::Error>
+    pub(super) fn deserialize<'de, D>(deserializer: D) -> Result<Value, D::Error>
     where
         D: Deserializer<'de>,
     {
@@ -88,7 +87,7 @@ mod args_serde_with {
         }
 
         impl<'de> Visitor<'de> for OptionVisitor {
-            type Value = Option<Vec<String>>;
+            type Value = Value;
 
             fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                 formatter.write_str("option")
diff --git a/crates/hyfetch/src/neofetch_util.rs b/crates/hyfetch/src/neofetch_util.rs
index e5cfcaf4..bec88a84 100644
--- a/crates/hyfetch/src/neofetch_util.rs
+++ b/crates/hyfetch/src/neofetch_util.rs
@@ -37,6 +37,7 @@ pub enum ColorAlignment {
     },
     Custom {
         #[serde(rename = "custom_colors")]
+        #[serde(deserialize_with = "crate::utils::index_map_serde::deserialize")]
         colors: IndexMap<NeofetchAsciiIndexedColor, PresetIndexedColor>,
     },
 }
@@ -185,8 +186,28 @@ impl ColorAlignment {
                     },
                 }
             },
-            Self::Custom { colors } => {
-                todo!()
+            Self::Custom {
+                colors: custom_colors,
+            } => {
+                let ColorProfile { colors } = color_profile.unique_colors();
+
+                // Apply colors
+                let asc = {
+                    let ac = NEOFETCH_COLORS_AC
+                        .get_or_init(|| AhoCorasick::new(NEOFETCH_COLOR_PATTERNS).unwrap());
+                    const N: usize = NEOFETCH_COLOR_PATTERNS.len();
+                    let mut replacements = vec![Cow::from(""); N];
+                    for (&ai, &pi) in custom_colors {
+                        let ai: u8 = ai.into();
+                        let pi: u8 = pi.into();
+                        replacements[ai as usize - 1] = colors[pi as usize]
+                            .to_ansi_string(color_mode, ForegroundBackground::Foreground)
+                            .into();
+                    }
+                    ac.replace_all(&asc, &replacements)
+                };
+
+                asc
             },
         };
 
diff --git a/crates/hyfetch/src/utils.rs b/crates/hyfetch/src/utils.rs
index b4571eda..ae061fe5 100644
--- a/crates/hyfetch/src/utils.rs
+++ b/crates/hyfetch/src/utils.rs
@@ -10,3 +10,97 @@ pub fn get_cache_path() -> Result<PathBuf> {
         .to_owned();
     Ok(path)
 }
+
+pub(crate) mod index_map_serde {
+    use std::fmt;
+    use std::fmt::Display;
+    use std::hash::Hash;
+    use std::marker::PhantomData;
+    use std::str::FromStr;
+
+    use indexmap::IndexMap;
+    use serde::de::{self, DeserializeSeed, MapAccess, Visitor};
+    use serde::{Deserialize, Deserializer};
+
+    pub(crate) fn deserialize<'de, D, K, V>(deserializer: D) -> Result<IndexMap<K, V>, D::Error>
+    where
+        D: Deserializer<'de>,
+        K: Eq + Hash + FromStr,
+        K::Err: Display,
+        V: Deserialize<'de>,
+    {
+        struct KeySeed<K> {
+            k: PhantomData<K>,
+        }
+
+        impl<'de, K> DeserializeSeed<'de> for KeySeed<K>
+        where
+            K: FromStr,
+            K::Err: Display,
+        {
+            type Value = K;
+
+            fn deserialize<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
+            where
+                D: Deserializer<'de>,
+            {
+                deserializer.deserialize_str(self)
+            }
+        }
+
+        impl<'de, K> Visitor<'de> for KeySeed<K>
+        where
+            K: FromStr,
+            K::Err: Display,
+        {
+            type Value = K;
+
+            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                formatter.write_str("a string")
+            }
+
+            fn visit_str<E>(self, s: &str) -> Result<Self::Value, E>
+            where
+                E: de::Error,
+            {
+                K::from_str(s).map_err(de::Error::custom)
+            }
+        }
+
+        struct MapVisitor<K, V> {
+            k: PhantomData<K>,
+            v: PhantomData<V>,
+        }
+
+        impl<'de, K, V> Visitor<'de> for MapVisitor<K, V>
+        where
+            K: Eq + Hash + FromStr,
+            K::Err: Display,
+            V: Deserialize<'de>,
+        {
+            type Value = IndexMap<K, V>;
+
+            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                formatter.write_str("a map")
+            }
+
+            fn visit_map<A>(self, mut input: A) -> Result<Self::Value, A::Error>
+            where
+                A: MapAccess<'de>,
+            {
+                let mut map = IndexMap::new();
+                while let Some((k, v)) =
+                    input.next_entry_seed(KeySeed { k: PhantomData }, PhantomData)?
+                {
+                    map.insert(k, v);
+                }
+                Ok(map)
+            }
+        }
+
+        deserializer.deserialize_map(MapVisitor {
+            k: PhantomData,
+            v: PhantomData,
+        })
+    }
+}