Skip to main content

leenfetch_core/core/
mod.rs

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