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                    && let Some(kind) = ModuleKind::from_field_name(field_name)
350                    && seen.insert(kind)
351                {
352                    required.push(kind);
353                }
354            }
355        }
356
357        if required.is_empty() {
358            return Data::default();
359        }
360
361        // Use rayon for parallel execution
362        let context = Arc::new(CollectContext::new(self.flags.clone()));
363        let modules: Vec<_> = required.into_iter().collect();
364
365        let results: Vec<Data> = modules
366            .into_par_iter()
367            .map(|kind| Self::collect_module_data(kind, context.clone()))
368            .collect();
369
370        let mut data = Data::default();
371        for update in results {
372            Self::merge_data(&mut data, update);
373        }
374
375        data
376    }
377
378    fn collect_module_data(kind: ModuleKind, context: Arc<CollectContext>) -> Data {
379        let mut data = Data::default();
380        let flags = &context.flags;
381
382        match kind {
383            ModuleKind::Titles => {
384                let (user, host) = get_titles(true);
385                data.username = Some(user);
386                data.hostname = Some(host);
387            }
388            ModuleKind::Os => {
389                data.os = Some(get_os());
390            }
391            ModuleKind::Distro => {
392                let display = DistroDisplay::from_str(&flags.distro_shorthand)
393                    .unwrap_or(DistroDisplay::NameModelVersionArch);
394                data.distro = Some(get_distro(display));
395            }
396            ModuleKind::Model => {
397                data.model = get_model();
398            }
399            ModuleKind::Kernel => {
400                data.kernel = get_kernel();
401            }
402            ModuleKind::OsAge => {
403                let os_age = get_os_age(
404                    OsAgeShorthand::from_str(&flags.os_age_shorthand)
405                        .unwrap_or(OsAgeShorthand::Tiny),
406                );
407                data.os_age = os_age;
408            }
409            ModuleKind::Uptime => {
410                let uptime = get_uptime(
411                    UptimeShorthand::from_str(&flags.uptime_shorthand)
412                        .unwrap_or(UptimeShorthand::Full),
413                );
414                data.uptime = uptime;
415            }
416            ModuleKind::Packages => {
417                let packages = get_packages(
418                    PackageShorthand::from_str(&flags.package_managers)
419                        .unwrap_or(PackageShorthand::On),
420                );
421                data.packages = packages;
422            }
423            ModuleKind::Shell => {
424                data.shell = get_shell(flags.shell_path, flags.shell_version);
425            }
426            ModuleKind::Wm => {
427                data.wm = context.get_wm();
428            }
429            ModuleKind::De => {
430                data.de = context.get_de();
431            }
432            ModuleKind::Cpu => {
433                data.cpu = get_cpu(
434                    flags.cpu_brand,
435                    flags.cpu_frequency,
436                    flags.cpu_cores,
437                    flags.cpu_temp != "off",
438                    flags.speed_shorthand,
439                    match flags.cpu_temp.as_str() {
440                        "C" => Some('C'),
441                        "F" => Some('F'),
442                        _ => None,
443                    },
444                );
445            }
446            ModuleKind::Gpu => {
447                data.gpu = Some(get_gpus());
448            }
449            ModuleKind::Memory => {
450                data.memory = get_memory(
451                    flags.memory_percent,
452                    MemoryUnit::from_str(flags.memory_unit.as_str()).unwrap_or(MemoryUnit::MiB),
453                );
454            }
455            ModuleKind::Disk => {
456                data.disk = get_disks(
457                    DiskSubtitle::from_str(flags.disk_subtitle.as_str())
458                        .unwrap_or(DiskSubtitle::Dir),
459                    DiskDisplay::from_str(flags.disk_display.as_str())
460                        .unwrap_or(DiskDisplay::InfoBar),
461                    None,
462                );
463            }
464            ModuleKind::Resolution => {
465                data.resolution = get_resolution();
466            }
467            ModuleKind::Theme => {
468                let de = context.get_de();
469                data.de = de.clone();
470                data.theme = get_theme(de.as_deref());
471            }
472            ModuleKind::Battery => {
473                let mode = BatteryDisplayMode::from_str(flags.battery_display.as_str())
474                    .unwrap_or(BatteryDisplayMode::BarInfo);
475                let batteries = get_battery(mode);
476                data.battery = Some(batteries);
477            }
478            ModuleKind::Song => {
479                data.song = get_song();
480            }
481            ModuleKind::Colors => {
482                let color_blocks = if flags.color_blocks.is_empty() {
483                    "●"
484                } else {
485                    flags.color_blocks.as_str()
486                };
487                let colors = get_terminal_color(color_blocks);
488                data.colors = Some(colors);
489            }
490        }
491
492        data
493    }
494
495    fn merge_data(target: &mut Data, update: Data) {
496        if let Some(username) = update.username {
497            target.username = Some(username);
498        }
499        if let Some(hostname) = update.hostname {
500            target.hostname = Some(hostname);
501        }
502        if let Some(os) = update.os {
503            target.os = Some(os);
504        }
505        if let Some(distro) = update.distro {
506            target.distro = Some(distro);
507        }
508        if let Some(model) = update.model {
509            target.model = Some(model);
510        }
511        if let Some(kernel) = update.kernel {
512            target.kernel = Some(kernel);
513        }
514        if let Some(os_age) = update.os_age {
515            target.os_age = Some(os_age);
516        }
517        if let Some(uptime) = update.uptime {
518            target.uptime = Some(uptime);
519        }
520        if let Some(packages) = update.packages {
521            target.packages = Some(packages);
522        }
523        if let Some(shell) = update.shell {
524            target.shell = Some(shell);
525        }
526        if let Some(wm) = update.wm {
527            target.wm = Some(wm);
528        }
529        if let Some(de) = update.de {
530            target.de = Some(de);
531        }
532        if let Some(cpu) = update.cpu {
533            target.cpu = Some(cpu);
534        }
535        if let Some(gpu) = update.gpu {
536            target.gpu = Some(gpu);
537        }
538        if let Some(memory) = update.memory {
539            target.memory = Some(memory);
540        }
541        if let Some(disk) = update.disk {
542            target.disk = Some(disk);
543        }
544        if let Some(resolution) = update.resolution {
545            target.resolution = Some(resolution);
546        }
547        if let Some(theme) = update.theme {
548            target.theme = Some(theme);
549        }
550        if let Some(battery) = update.battery {
551            target.battery = Some(battery);
552        }
553        if let Some(song) = update.song {
554            target.song = Some(song);
555        }
556        if let Some(colors) = update.colors {
557            target.colors = Some(colors);
558        }
559    }
560
561    fn push_unknown(label: &str, output: &mut String) {
562        output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, "Unknown").as_str());
563    }
564
565    pub fn get_ascii_and_colors(&self) -> (String, HashMap<&str, &str>) {
566        self.get_ascii_and_colors_for_distro(None)
567    }
568
569    pub fn get_ascii_and_colors_for_distro(
570        &self,
571        distro_override: Option<&str>,
572    ) -> (String, HashMap<&str, &str>) {
573        let custom_logo_path = {
574            let path = self.flags.custom_logo_path.trim();
575            if path.is_empty() { None } else { Some(path) }
576        };
577
578        let ascii_color_value = {
579            let value = self.flags.ascii_colors.trim();
580            if value.is_empty() { "distro" } else { value }
581        };
582
583        let ascii_distro_value = {
584            let value = self.flags.ascii_distro.trim();
585            if value.is_empty() { "distro" } else { value }
586        };
587
588        let resolved_distro = match ascii_distro_value {
589            "auto" => distro_override
590                .map(|s| s.to_string())
591                .unwrap_or_else(|| get_distro(DistroDisplay::Name)),
592            "auto_small" => {
593                let base = distro_override
594                    .map(|s| s.to_string())
595                    .unwrap_or_else(|| get_distro(DistroDisplay::Name));
596                format!("{}_small", base)
597            }
598            other => other.to_string(),
599        };
600
601        // Load ASCII Art
602        let raw_ascii_art = custom_logo_path
603            .map(get_custom_ascii)
604            .unwrap_or_else(|| get_ascii_and_colors(&resolved_distro));
605
606        // Load Colors
607        let distro_colors = if &resolved_distro == "off" {
608            get_distro_colors(&get_distro(DistroDisplay::Name))
609        } else {
610            match ascii_color_value {
611                "distro" => get_distro_colors(&resolved_distro),
612                other => get_custom_colors_order(other),
613            }
614        };
615
616        (raw_ascii_art, distro_colors)
617    }
618
619    /// If the given `data` is `Some`, it will be added to `output` with the given `label`.
620    /// If `data` is `None`, it will add a line to `output` with the label and the value "Unknown".
621    fn is_some_add_to_output(label: &str, data: &Option<String>, output: &mut String) {
622        match data {
623            Some(d) => {
624                output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, d).as_str());
625            }
626            None => {
627                output.push_str(format!("${{c1}}{} ${{reset}}{}\n", label, "Unknown").as_str());
628            }
629        }
630    }
631}
632
633impl Default for Core {
634    fn default() -> Self {
635        Self::new()
636    }
637}