Skip to main content

leenfetch_core/core/
mod.rs

1mod data;
2
3pub use data::Data;
4
5use std::{collections::HashMap, str::FromStr, sync::Arc};
6
7use once_cell::sync::OnceCell;
8use rayon::prelude::*;
9
10use crate::{
11    config::{self, settings},
12    modules::{
13        desktop::{de::get_de, resolution::get_resolution, theme::get_theme, wm::get_wm},
14        enums::{
15            BatteryDisplayMode, DiskSubtitle, DistroDisplay, MemoryUnit, OsAgeShorthand,
16            PackageShorthand, UptimeShorthand,
17        },
18        info::{
19            battery::get_battery, cpu::get_cpu, disk::get_disks, gpu::get_gpus, memory::get_memory,
20            os_age::get_os_age, uptime::get_uptime,
21        },
22        packages::get_packages,
23        shell::get_shell,
24        song::get_song,
25        system::{distro::get_distro, kernel::get_kernel, model::get_model, os::get_os},
26        title::get_titles,
27        utils::{
28            get_ascii_and_colors, get_custom_ascii, get_custom_colors_order, get_distro_colors,
29            get_terminal_color,
30        },
31    },
32};
33
34#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
35enum ModuleKind {
36    Titles,
37    Os,
38    Distro,
39    Model,
40    Kernel,
41    OsAge,
42    Uptime,
43    Packages,
44    Shell,
45    Wm,
46    De,
47    Cpu,
48    Gpu,
49    Memory,
50    Disk,
51    Resolution,
52    Theme,
53    Battery,
54    Song,
55    Colors,
56}
57
58impl ModuleKind {
59    fn from_field_name(name: &str) -> Option<Self> {
60        let normalized = name.trim().to_ascii_lowercase().replace('-', "_");
61        match normalized.as_str() {
62            "titles" => Some(Self::Titles),
63            "os" => Some(Self::Os),
64            "distro" => Some(Self::Distro),
65            "model" => Some(Self::Model),
66            "kernel" => Some(Self::Kernel),
67            "os_age" => Some(Self::OsAge),
68            "uptime" => Some(Self::Uptime),
69            "packages" => Some(Self::Packages),
70            "shell" => Some(Self::Shell),
71            "wm" => Some(Self::Wm),
72            "de" => Some(Self::De),
73            "cpu" => Some(Self::Cpu),
74            "gpu" => Some(Self::Gpu),
75            "memory" => Some(Self::Memory),
76            "disk" => Some(Self::Disk),
77            "resolution" => Some(Self::Resolution),
78            "theme" => Some(Self::Theme),
79            "battery" => Some(Self::Battery),
80            "song" => Some(Self::Song),
81            "colors" => Some(Self::Colors),
82            _ => None,
83        }
84    }
85}
86
87struct CollectContext {
88    flags: settings::Flags,
89    wm: OnceCell<Option<String>>,
90    de: OnceCell<Option<String>>,
91}
92
93impl CollectContext {
94    fn new(flags: settings::Flags) -> Self {
95        Self {
96            flags,
97            wm: OnceCell::new(),
98            de: OnceCell::new(),
99        }
100    }
101
102    fn get_wm(&self) -> Option<String> {
103        self.wm.get_or_init(get_wm).clone()
104    }
105
106    fn get_de(&self) -> Option<String> {
107        self.de
108            .get_or_init(|| {
109                let wm = self.get_wm();
110                get_de(self.flags.de_version, wm.as_deref())
111            })
112            .clone()
113    }
114}
115
116pub struct Core {
117    flags: settings::Flags,
118    layout: Vec<settings::LayoutItem>,
119}
120
121impl Core {
122    /// Creates a new instance of the `Core` struct.
123    ///
124    /// This function reads the configuration file and populates the `Core` struct with
125    /// the configuration and default system information.
126    pub fn new() -> Self {
127        let flags = config::load_flags();
128        let layout = config::load_print_layout();
129        Self::new_with(flags, layout)
130    }
131
132    /// Creates a new `Core` from explicit flags and layout values.
133    pub fn new_with(flags: settings::Flags, layout: Vec<settings::LayoutItem>) -> Self {
134        Self { flags, layout }
135    }
136
137    /// Builds the final colorized layout output using the loaded configuration.
138    ///
139    /// Each entry in the layout is resolved against the configured flags. Module data is collected
140    /// in parallel and rendered in the configured order so that expensive lookups overlap without
141    /// changing how the output is composed.
142    ///
143    /// # Returns
144    ///
145    /// A single colorized `String` containing the assembled module output.
146    pub fn get_info_layout(&self) -> String {
147        let data = self.collect_data_parallel();
148        self.render_layout(&data)
149    }
150
151    /// Renders the layout using pre-collected data.
152    pub fn render_layout(&self, data: &Data) -> String {
153        let mut final_output = String::new();
154
155        for item in &self.layout {
156            match item {
157                settings::LayoutItem::Break(value) => {
158                    if value.eq_ignore_ascii_case("break") {
159                        final_output.push('\n');
160                    } else {
161                        final_output.push_str(value);
162                        if !value.ends_with('\n') {
163                            final_output.push('\n');
164                        }
165                    }
166                }
167                settings::LayoutItem::Module(module) => {
168                    let Some(raw_field_name) = module.field_name() else {
169                        continue;
170                    };
171
172                    if module.is_custom() {
173                        if let Some(text) = module.format.as_ref().or(module.text.as_ref()) {
174                            final_output.push_str(text);
175                            if !text.ends_with('\n') {
176                                final_output.push('\n');
177                            }
178                        }
179                        continue;
180                    }
181
182                    let field_name = raw_field_name.to_ascii_lowercase();
183                    let label_storage = module
184                        .label()
185                        .map(|value| value.to_string())
186                        .unwrap_or_else(|| field_name.clone());
187                    let label = label_storage.as_str();
188
189                    match ModuleKind::from_field_name(&field_name) {
190                        Some(ModuleKind::Titles) => {
191                            let username = data.username.as_deref().unwrap_or("Unknown");
192                            let hostname = data.hostname.as_deref().unwrap_or("Unknown");
193                            let titles_line = format!(
194                                "${{c1}}{}${{reset}} {}${{c1}}@${{reset}}{}${{reset}}\n",
195                                label, username, hostname,
196                            );
197                            final_output.push_str(&titles_line);
198                        }
199                        Some(ModuleKind::Os) => {
200                            Self::is_some_add_to_output(label, &data.os, &mut final_output);
201                        }
202                        Some(ModuleKind::Distro) => {
203                            Self::is_some_add_to_output(label, &data.distro, &mut final_output);
204                        }
205                        Some(ModuleKind::Model) => {
206                            Self::is_some_add_to_output(label, &data.model, &mut final_output);
207                        }
208                        Some(ModuleKind::Kernel) => {
209                            Self::is_some_add_to_output(label, &data.kernel, &mut final_output);
210                        }
211                        Some(ModuleKind::OsAge) => {
212                            Self::is_some_add_to_output(label, &data.os_age, &mut final_output);
213                        }
214                        Some(ModuleKind::Uptime) => {
215                            Self::is_some_add_to_output(label, &data.uptime, &mut final_output);
216                        }
217                        Some(ModuleKind::Packages) => {
218                            Self::is_some_add_to_output(label, &data.packages, &mut final_output);
219                        }
220                        Some(ModuleKind::Shell) => {
221                            Self::is_some_add_to_output(label, &data.shell, &mut final_output);
222                        }
223                        Some(ModuleKind::Wm) => {
224                            Self::is_some_add_to_output(label, &data.wm, &mut final_output);
225                        }
226                        Some(ModuleKind::De) => {
227                            Self::is_some_add_to_output(label, &data.de, &mut final_output);
228                        }
229                        Some(ModuleKind::Cpu) => {
230                            Self::is_some_add_to_output(label, &data.cpu, &mut final_output);
231                        }
232                        Some(ModuleKind::Gpu) => match data.gpu.as_ref() {
233                            Some(gpus) if gpus.is_empty() => {
234                                let line =
235                                    format!("${{c1}}{} ${{reset}}{}\n", label, "No GPU found");
236                                final_output.push_str(&line);
237                            }
238                            Some(gpus) if gpus.len() == 1 => {
239                                let line = format!("${{c1}}{} ${{reset}}{}\n", label, gpus[0]);
240                                final_output.push_str(&line);
241                            }
242                            Some(gpus) => {
243                                for gpu in gpus {
244                                    let line = format!("${{c1}}{} ${{reset}}{}\n", label, gpu);
245                                    final_output.push_str(&line);
246                                }
247                            }
248                            None => Self::push_unknown(label, &mut final_output),
249                        },
250                        Some(ModuleKind::Memory) => {
251                            Self::is_some_add_to_output(label, &data.memory, &mut final_output);
252                        }
253                        Some(ModuleKind::Disk) => match data.disk.as_ref() {
254                            Some(disks) => {
255                                if disks.is_empty() {
256                                    let line = format!(
257                                        "${{c1}}{} ${{reset}}{}\n",
258                                        label, "No disks found"
259                                    );
260                                    final_output.push_str(&line);
261                                } else {
262                                    for (name, summary) in disks {
263                                        let line = if name.is_empty() {
264                                            format!("${{c1}}{} ${{reset}}{}\n", label, summary)
265                                        } else {
266                                            format!(
267                                                "${{c1}}{} {} ${{reset}}{}\n",
268                                                label, name, summary
269                                            )
270                                        };
271                                        final_output.push_str(&line);
272                                    }
273                                }
274                            }
275                            None => {
276                                let line =
277                                    format!("${{c1}}{} ${{reset}}{}\n", label, "No disks found");
278                                final_output.push_str(&line);
279                            }
280                        },
281                        Some(ModuleKind::Resolution) => {
282                            Self::is_some_add_to_output(label, &data.resolution, &mut final_output);
283                        }
284                        Some(ModuleKind::Theme) => {
285                            Self::is_some_add_to_output(label, &data.theme, &mut final_output);
286                        }
287                        Some(ModuleKind::Battery) => match data.battery.as_ref() {
288                            Some(batteries) if batteries.is_empty() => {
289                                let line =
290                                    format!("${{c1}}{} ${{reset}}{}\n", label, "No Battery found");
291                                final_output.push_str(&line);
292                            }
293                            Some(batteries) if batteries.len() == 1 => {
294                                let line = format!("${{c1}}{} ${{reset}}{}\n", label, batteries[0]);
295                                final_output.push_str(&line);
296                            }
297                            Some(batteries) => {
298                                for (index, battery) in batteries.iter().enumerate() {
299                                    let line = format!(
300                                        "${{c1}}{} {}: ${{reset}}{}\n",
301                                        label, index, battery
302                                    );
303                                    final_output.push_str(&line);
304                                }
305                            }
306                            None => {
307                                let line =
308                                    format!("${{c1}}{} ${{reset}}{}\n", label, "No Battery found");
309                                final_output.push_str(&line);
310                            }
311                        },
312                        Some(ModuleKind::Song) => {
313                            if let Some(music) = data.song.as_ref() {
314                                let line = format!(
315                                    "${{c1}}Playing${{reset}}\n    {}\n    {}\n",
316                                    music.title, music.artist
317                                );
318                                final_output.push_str(&line);
319                            }
320                        }
321                        Some(ModuleKind::Colors) => {
322                            Self::is_some_add_to_output(label, &data.colors, &mut final_output);
323                        }
324                        None => {
325                            let fallback_line =
326                                format!("${{c1}}{} ${{reset}}{}\n", label, field_name);
327                            final_output.push_str(&fallback_line);
328                        }
329                    }
330                }
331            }
332        }
333
334        final_output
335    }
336
337    /// Collects data for all modules referenced in the configured layout.
338    pub fn collect_data(&self) -> Data {
339        self.collect_data_parallel()
340    }
341
342    fn collect_data_parallel(&self) -> Data {
343        let mut required = Vec::new();
344        let mut seen = std::collections::HashSet::new();
345
346        for item in &self.layout {
347            if let settings::LayoutItem::Module(module) = item {
348                if module.is_custom() {
349                    continue;
350                }
351
352                if let Some(field_name) = module.field_name()
353                    && let Some(kind) = ModuleKind::from_field_name(field_name)
354                    && seen.insert(kind)
355                {
356                    required.push(kind);
357                }
358            }
359        }
360
361        if required.is_empty() {
362            return Data::default();
363        }
364
365        // Use rayon for parallel execution
366        let context = Arc::new(CollectContext::new(self.flags.clone()));
367        let modules: Vec<_> = required.into_iter().collect();
368
369        let results: Vec<Data> = modules
370            .into_par_iter()
371            .map(|kind| Self::collect_module_data(kind, context.clone()))
372            .collect();
373
374        let mut data = Data::default();
375        for update in results {
376            Self::merge_data(&mut data, update);
377        }
378
379        data
380    }
381
382    fn collect_module_data(kind: ModuleKind, context: Arc<CollectContext>) -> Data {
383        let mut data = Data::default();
384        let flags = &context.flags;
385
386        match kind {
387            ModuleKind::Titles => {
388                let (user, host) = get_titles(true);
389                data.username = Some(user);
390                data.hostname = Some(host);
391            }
392            ModuleKind::Os => {
393                data.os = Some(get_os());
394            }
395            ModuleKind::Distro => {
396                let display = DistroDisplay::from_str(&flags.distro_shorthand)
397                    .unwrap_or(DistroDisplay::NameModelVersionArch);
398                data.distro = Some(get_distro(display));
399            }
400            ModuleKind::Model => {
401                data.model = get_model();
402            }
403            ModuleKind::Kernel => {
404                data.kernel = get_kernel();
405            }
406            ModuleKind::OsAge => {
407                let os_age = get_os_age(
408                    OsAgeShorthand::from_str(&flags.os_age_shorthand)
409                        .unwrap_or(OsAgeShorthand::Tiny),
410                );
411                data.os_age = os_age;
412            }
413            ModuleKind::Uptime => {
414                let uptime = get_uptime(
415                    UptimeShorthand::from_str(&flags.uptime_shorthand)
416                        .unwrap_or(UptimeShorthand::Full),
417                );
418                data.uptime = uptime;
419            }
420            ModuleKind::Packages => {
421                let packages = get_packages(
422                    PackageShorthand::from_str(&flags.package_managers)
423                        .unwrap_or(PackageShorthand::On),
424                );
425                data.packages = packages;
426            }
427            ModuleKind::Shell => {
428                data.shell = get_shell(flags.shell_path, flags.shell_version);
429            }
430            ModuleKind::Wm => {
431                data.wm = context.get_wm();
432            }
433            ModuleKind::De => {
434                data.de = context.get_de();
435            }
436            ModuleKind::Cpu => {
437                data.cpu = get_cpu(
438                    flags.cpu_brand,
439                    flags.cpu_frequency,
440                    flags.cpu_cores,
441                    flags.cpu_temp != "off",
442                    flags.speed_shorthand,
443                    match flags.cpu_temp.as_str() {
444                        "C" => Some('C'),
445                        "F" => Some('F'),
446                        _ => None,
447                    },
448                );
449            }
450            ModuleKind::Gpu => {
451                data.gpu = Some(get_gpus());
452            }
453            ModuleKind::Memory => {
454                data.memory = get_memory(
455                    flags.memory_percent,
456                    MemoryUnit::from_str(flags.memory_unit.as_str()).unwrap_or(MemoryUnit::MiB),
457                );
458            }
459            ModuleKind::Disk => {
460                data.disk = get_disks(
461                    DiskSubtitle::from_str(flags.disk_subtitle.as_str())
462                        .unwrap_or(DiskSubtitle::Dir),
463                    flags.disk_display.as_str(),
464                    None,
465                );
466            }
467            ModuleKind::Resolution => {
468                data.resolution = get_resolution();
469            }
470            ModuleKind::Theme => {
471                let de = context.get_de();
472                data.de = de.clone();
473                data.theme = get_theme(de.as_deref());
474            }
475            ModuleKind::Battery => {
476                let mode = BatteryDisplayMode::from_str(flags.battery_display.as_str())
477                    .unwrap_or(BatteryDisplayMode::BarInfo);
478                let batteries = get_battery(mode);
479                data.battery = Some(batteries);
480            }
481            ModuleKind::Song => {
482                data.song = get_song();
483            }
484            ModuleKind::Colors => {
485                let color_blocks = if flags.color_blocks.is_empty() {
486                    "●"
487                } else {
488                    flags.color_blocks.as_str()
489                };
490                let colors = get_terminal_color(color_blocks);
491                data.colors = Some(colors);
492            }
493        }
494
495        data
496    }
497
498    fn merge_data(target: &mut Data, update: Data) {
499        if let Some(username) = update.username {
500            target.username = Some(username);
501        }
502        if let Some(hostname) = update.hostname {
503            target.hostname = Some(hostname);
504        }
505        if let Some(os) = update.os {
506            target.os = Some(os);
507        }
508        if let Some(distro) = update.distro {
509            target.distro = Some(distro);
510        }
511        if let Some(model) = update.model {
512            target.model = Some(model);
513        }
514        if let Some(kernel) = update.kernel {
515            target.kernel = Some(kernel);
516        }
517        if let Some(os_age) = update.os_age {
518            target.os_age = Some(os_age);
519        }
520        if let Some(uptime) = update.uptime {
521            target.uptime = Some(uptime);
522        }
523        if let Some(packages) = update.packages {
524            target.packages = Some(packages);
525        }
526        if let Some(shell) = update.shell {
527            target.shell = Some(shell);
528        }
529        if let Some(wm) = update.wm {
530            target.wm = Some(wm);
531        }
532        if let Some(de) = update.de {
533            target.de = Some(de);
534        }
535        if let Some(cpu) = update.cpu {
536            target.cpu = Some(cpu);
537        }
538        if let Some(gpu) = update.gpu {
539            target.gpu = Some(gpu);
540        }
541        if let Some(memory) = update.memory {
542            target.memory = Some(memory);
543        }
544        if let Some(disk) = update.disk {
545            target.disk = Some(disk);
546        }
547        if let Some(resolution) = update.resolution {
548            target.resolution = Some(resolution);
549        }
550        if let Some(theme) = update.theme {
551            target.theme = Some(theme);
552        }
553        if let Some(battery) = update.battery {
554            target.battery = Some(battery);
555        }
556        if let Some(song) = update.song {
557            target.song = Some(song);
558        }
559        if let Some(colors) = update.colors {
560            target.colors = Some(colors);
561        }
562    }
563
564    fn push_unknown(label: &str, output: &mut String) {
565        output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, "Unknown").as_str());
566    }
567
568    pub fn get_ascii_and_colors(&self) -> (String, HashMap<&str, &str>) {
569        self.get_ascii_and_colors_for_distro(None)
570    }
571
572    pub fn get_ascii_and_colors_for_distro(
573        &self,
574        distro_override: Option<&str>,
575    ) -> (String, HashMap<&str, &str>) {
576        let custom_logo_path = {
577            let path = self.flags.custom_logo_path.trim();
578            if path.is_empty() { None } else { Some(path) }
579        };
580
581        let ascii_color_value = {
582            let value = self.flags.ascii_colors.trim();
583            if value.is_empty() { "distro" } else { value }
584        };
585
586        let ascii_distro_value = {
587            let value = self.flags.ascii_distro.trim();
588            if value.is_empty() { "distro" } else { value }
589        };
590
591        let resolved_distro = match ascii_distro_value {
592            "auto" => distro_override
593                .map(|s| s.to_string())
594                .unwrap_or_else(|| get_distro(DistroDisplay::Name)),
595            "auto_small" => {
596                let base = distro_override
597                    .map(|s| s.to_string())
598                    .unwrap_or_else(|| get_distro(DistroDisplay::Name));
599                format!("{}_small", base)
600            }
601            other => other.to_string(),
602        };
603
604        // Load ASCII Art
605        let raw_ascii_art = custom_logo_path
606            .map(get_custom_ascii)
607            .unwrap_or_else(|| get_ascii_and_colors(&resolved_distro));
608
609        // Load Colors
610        let distro_colors = if &resolved_distro == "off" {
611            get_distro_colors(&get_distro(DistroDisplay::Name))
612        } else {
613            match ascii_color_value {
614                "distro" => get_distro_colors(&resolved_distro),
615                other => get_custom_colors_order(other),
616            }
617        };
618
619        (raw_ascii_art, distro_colors)
620    }
621
622    /// If the given `data` is `Some`, it will be added to `output` with the given `label`.
623    /// If `data` is `None`, it will add a line to `output` with the label and the value "Unknown".
624    fn is_some_add_to_output(label: &str, data: &Option<String>, output: &mut String) {
625        match data {
626            Some(d) => {
627                output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, d).as_str());
628            }
629            None => {
630                output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, "Unknown").as_str());
631            }
632        }
633    }
634}
635
636impl Default for Core {
637    fn default() -> Self {
638        Self::new()
639    }
640}