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, DiskDisplay, DiskSubtitle, DistroDisplay, MemoryUnit,
16            OsAgeShorthand, 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 = format!(
264                                            "${{c1}}{} {} ${{reset}}{}\n",
265                                            label, name, summary
266                                        );
267                                        final_output.push_str(&line);
268                                    }
269                                }
270                            }
271                            None => {
272                                let line =
273                                    format!("${{c1}}{} ${{reset}}{}\n", label, "No disks found");
274                                final_output.push_str(&line);
275                            }
276                        },
277                        Some(ModuleKind::Resolution) => {
278                            Self::is_some_add_to_output(label, &data.resolution, &mut final_output);
279                        }
280                        Some(ModuleKind::Theme) => {
281                            Self::is_some_add_to_output(label, &data.theme, &mut final_output);
282                        }
283                        Some(ModuleKind::Battery) => match data.battery.as_ref() {
284                            Some(batteries) if batteries.is_empty() => {
285                                let line =
286                                    format!("${{c1}}{} ${{reset}}{}\n", label, "No Battery found");
287                                final_output.push_str(&line);
288                            }
289                            Some(batteries) if batteries.len() == 1 => {
290                                let line = format!("${{c1}}{} ${{reset}}{}\n", label, batteries[0]);
291                                final_output.push_str(&line);
292                            }
293                            Some(batteries) => {
294                                for (index, battery) in batteries.iter().enumerate() {
295                                    let line = format!(
296                                        "${{c1}}{} {}: ${{reset}}{}\n",
297                                        label, index, battery
298                                    );
299                                    final_output.push_str(&line);
300                                }
301                            }
302                            None => {
303                                let line =
304                                    format!("${{c1}}{} ${{reset}}{}\n", label, "No Battery found");
305                                final_output.push_str(&line);
306                            }
307                        },
308                        Some(ModuleKind::Song) => {
309                            if let Some(music) = data.song.as_ref() {
310                                let line = format!(
311                                    "${{c1}}Playing${{reset}}\n    {}\n    {}\n",
312                                    music.title, music.artist
313                                );
314                                final_output.push_str(&line);
315                            }
316                        }
317                        Some(ModuleKind::Colors) => {
318                            Self::is_some_add_to_output(label, &data.colors, &mut final_output);
319                        }
320                        None => {
321                            let fallback_line =
322                                format!("${{c1}}{} ${{reset}}{}\n", label, field_name);
323                            final_output.push_str(&fallback_line);
324                        }
325                    }
326                }
327            }
328        }
329
330        final_output
331    }
332
333    /// Collects data for all modules referenced in the configured layout.
334    pub fn collect_data(&self) -> Data {
335        self.collect_data_parallel()
336    }
337
338    fn collect_data_parallel(&self) -> Data {
339        let mut required = Vec::new();
340        let mut seen = std::collections::HashSet::new();
341
342        for item in &self.layout {
343            if let settings::LayoutItem::Module(module) = item {
344                if module.is_custom() {
345                    continue;
346                }
347
348                if let Some(field_name) = module.field_name() {
349                    if let Some(kind) = ModuleKind::from_field_name(field_name) {
350                        if seen.insert(kind) {
351                            required.push(kind);
352                        }
353                    }
354                }
355            }
356        }
357
358        if required.is_empty() {
359            return Data::default();
360        }
361
362        // Use rayon for parallel execution
363        let context = Arc::new(CollectContext::new(self.flags.clone()));
364        let modules: Vec<_> = required.into_iter().collect();
365
366        let results: Vec<Data> = modules
367            .into_par_iter()
368            .map(|kind| Self::collect_module_data(kind, context.clone()))
369            .collect();
370
371        let mut data = Data::default();
372        for update in results {
373            Self::merge_data(&mut data, update);
374        }
375
376        data
377    }
378
379    fn collect_module_data(kind: ModuleKind, context: Arc<CollectContext>) -> Data {
380        let mut data = Data::default();
381        let flags = &context.flags;
382
383        match kind {
384            ModuleKind::Titles => {
385                let (user, host) = get_titles(true);
386                data.username = Some(user);
387                data.hostname = Some(host);
388            }
389            ModuleKind::Os => {
390                data.os = Some(get_os());
391            }
392            ModuleKind::Distro => {
393                let display = DistroDisplay::from_str(&flags.distro_shorthand)
394                    .unwrap_or(DistroDisplay::NameModelVersionArch);
395                data.distro = Some(get_distro(display));
396            }
397            ModuleKind::Model => {
398                data.model = get_model();
399            }
400            ModuleKind::Kernel => {
401                data.kernel = get_kernel();
402            }
403            ModuleKind::OsAge => {
404                let os_age = get_os_age(
405                    OsAgeShorthand::from_str(&flags.os_age_shorthand)
406                        .unwrap_or(OsAgeShorthand::Tiny),
407                );
408                data.os_age = os_age;
409            }
410            ModuleKind::Uptime => {
411                let uptime = get_uptime(
412                    UptimeShorthand::from_str(&flags.uptime_shorthand)
413                        .unwrap_or(UptimeShorthand::Full),
414                );
415                data.uptime = uptime;
416            }
417            ModuleKind::Packages => {
418                let packages = get_packages(
419                    PackageShorthand::from_str(&flags.package_managers)
420                        .unwrap_or(PackageShorthand::On),
421                );
422                data.packages = packages;
423            }
424            ModuleKind::Shell => {
425                data.shell = get_shell(flags.shell_path, flags.shell_version);
426            }
427            ModuleKind::Wm => {
428                data.wm = context.get_wm();
429            }
430            ModuleKind::De => {
431                data.de = context.get_de();
432            }
433            ModuleKind::Cpu => {
434                data.cpu = get_cpu(
435                    flags.cpu_brand,
436                    flags.cpu_frequency,
437                    flags.cpu_cores,
438                    flags.cpu_temp != "off",
439                    flags.speed_shorthand,
440                    match flags.cpu_temp.as_str() {
441                        "C" => Some('C'),
442                        "F" => Some('F'),
443                        _ => None,
444                    },
445                );
446            }
447            ModuleKind::Gpu => {
448                data.gpu = Some(get_gpus());
449            }
450            ModuleKind::Memory => {
451                data.memory = get_memory(
452                    flags.memory_percent,
453                    MemoryUnit::from_str(flags.memory_unit.as_str()).unwrap_or(MemoryUnit::MiB),
454                );
455            }
456            ModuleKind::Disk => {
457                data.disk = get_disks(
458                    DiskSubtitle::from_str(flags.disk_subtitle.as_str())
459                        .unwrap_or(DiskSubtitle::Dir),
460                    DiskDisplay::from_str(flags.disk_display.as_str())
461                        .unwrap_or(DiskDisplay::InfoBar),
462                    None,
463                );
464            }
465            ModuleKind::Resolution => {
466                data.resolution = get_resolution();
467            }
468            ModuleKind::Theme => {
469                let de = context.get_de();
470                data.de = de.clone();
471                data.theme = get_theme(de.as_deref());
472            }
473            ModuleKind::Battery => {
474                let mode = BatteryDisplayMode::from_str(flags.battery_display.as_str())
475                    .unwrap_or(BatteryDisplayMode::BarInfo);
476                let batteries = get_battery(mode);
477                data.battery = Some(batteries);
478            }
479            ModuleKind::Song => {
480                data.song = get_song();
481            }
482            ModuleKind::Colors => {
483                let color_blocks = if flags.color_blocks.is_empty() {
484                    "●"
485                } else {
486                    flags.color_blocks.as_str()
487                };
488                let colors = get_terminal_color(color_blocks);
489                data.colors = Some(colors);
490            }
491        }
492
493        data
494    }
495
496    fn merge_data(target: &mut Data, update: Data) {
497        if let Some(username) = update.username {
498            target.username = Some(username);
499        }
500        if let Some(hostname) = update.hostname {
501            target.hostname = Some(hostname);
502        }
503        if let Some(os) = update.os {
504            target.os = Some(os);
505        }
506        if let Some(distro) = update.distro {
507            target.distro = Some(distro);
508        }
509        if let Some(model) = update.model {
510            target.model = Some(model);
511        }
512        if let Some(kernel) = update.kernel {
513            target.kernel = Some(kernel);
514        }
515        if let Some(os_age) = update.os_age {
516            target.os_age = Some(os_age);
517        }
518        if let Some(uptime) = update.uptime {
519            target.uptime = Some(uptime);
520        }
521        if let Some(packages) = update.packages {
522            target.packages = Some(packages);
523        }
524        if let Some(shell) = update.shell {
525            target.shell = Some(shell);
526        }
527        if let Some(wm) = update.wm {
528            target.wm = Some(wm);
529        }
530        if let Some(de) = update.de {
531            target.de = Some(de);
532        }
533        if let Some(cpu) = update.cpu {
534            target.cpu = Some(cpu);
535        }
536        if let Some(gpu) = update.gpu {
537            target.gpu = Some(gpu);
538        }
539        if let Some(memory) = update.memory {
540            target.memory = Some(memory);
541        }
542        if let Some(disk) = update.disk {
543            target.disk = Some(disk);
544        }
545        if let Some(resolution) = update.resolution {
546            target.resolution = Some(resolution);
547        }
548        if let Some(theme) = update.theme {
549            target.theme = Some(theme);
550        }
551        if let Some(battery) = update.battery {
552            target.battery = Some(battery);
553        }
554        if let Some(song) = update.song {
555            target.song = Some(song);
556        }
557        if let Some(colors) = update.colors {
558            target.colors = Some(colors);
559        }
560    }
561
562    fn push_unknown(label: &str, output: &mut String) {
563        output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, "Unknown").as_str());
564    }
565
566    pub fn get_ascii_and_colors(&self) -> (String, HashMap<&str, &str>) {
567        self.get_ascii_and_colors_for_distro(None)
568    }
569
570    pub fn get_ascii_and_colors_for_distro(
571        &self,
572        distro_override: Option<&str>,
573    ) -> (String, HashMap<&str, &str>) {
574        let custom_ascii_path = {
575            let path = self.flags.custom_ascii_path.trim();
576            if path.is_empty() {
577                None
578            } else {
579                Some(path)
580            }
581        };
582
583        let ascii_color_value = {
584            let value = self.flags.ascii_colors.trim();
585            if value.is_empty() {
586                "distro"
587            } else {
588                value
589            }
590        };
591
592        let ascii_distro_value = {
593            let value = self.flags.ascii_distro.trim();
594            if value.is_empty() {
595                "distro"
596            } else {
597                value
598            }
599        };
600
601        let resolved_distro = match ascii_distro_value {
602            "auto" => distro_override
603                .map(|s| s.to_string())
604                .unwrap_or_else(|| get_distro(DistroDisplay::Name)),
605            "auto_small" => {
606                let base = distro_override
607                    .map(|s| s.to_string())
608                    .unwrap_or_else(|| get_distro(DistroDisplay::Name));
609                format!("{}_small", base)
610            }
611            other => other.to_string(),
612        };
613
614        // Load ASCII Art
615        let raw_ascii_art = custom_ascii_path
616            .map(get_custom_ascii)
617            .unwrap_or_else(|| get_ascii_and_colors(&resolved_distro));
618
619        // Load Colors
620        let distro_colors = if &resolved_distro == "off" {
621            get_distro_colors(&get_distro(DistroDisplay::Name))
622        } else {
623            match ascii_color_value {
624                "distro" => get_distro_colors(&resolved_distro),
625                other => get_custom_colors_order(other),
626            }
627        };
628
629        (raw_ascii_art, distro_colors)
630    }
631
632    /// If the given `data` is `Some`, it will be added to `output` with the given `label`.
633    /// If `data` is `None`, it will add a line to `output` with the label and the value "Unknown".
634    fn is_some_add_to_output(label: &str, data: &Option<String>, output: &mut String) {
635        match data {
636            Some(d) => {
637                output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, d).as_str());
638            }
639            None => {
640                output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, "Unknown").as_str());
641            }
642        }
643    }
644}