#!/usr/bin/env python3

import filecmp
import os
from pathlib import Path
import pty
import re
import sys
import subprocess
import argparse
import json
from enum import Enum, auto
import time
from typing import Callable, Optional, Dict, NamedTuple, Set, Tuple, TypeAlias
import typing
from dataclasses import dataclass

# -- ANSI color codes --
COLORS = {
    "CRITICAL": "\033[91m", "HIGH": "\033[31m", "LOW": "\033[33m",
    "INFO": "\033[34m", "RESET": "\033[0m", "GREEN": "\033[32m",
}

# -- Core data structures --
Fixer = typing.NewType("Fixer", None)

class Severity(Enum):
    """Defines the 4 categories of issues and 3 additional states."""
    NO_ISSUE = auto()
    # Not applicable does not necessarily mean that the check could not be run
    # Sometimes the check can not be run, but it is still NO_ISSUE because NOT_APPLICABLE
    # is supposed to mean that there is missing software or another issue preventing it from running.
    # For example, missing thermal_zones are not the fault of the user and "unsolvable", therefore NO_ISSUE
    NOT_APPLICABLE = auto()
    ERROR = auto()
    INFO = auto()
    LOW = auto()
    HIGH = auto()
    CRITICAL = auto()

printable_severities = [
    Severity.INFO, Severity.LOW, Severity.HIGH, Severity.CRITICAL
]

@dataclass
class PackageInfo():
    installed: bool; provided: bool

class Issue:
    """
    Represents a specific problem.

    :ivar description: A human-readable string explaining the issue.
    :ivar details: A dictionary to pass context from check to fix.
    :ivar fix: A Fixer object that can be used to fix the issue.
    """
    def __init__(self,
                 description: str,
                 details: Optional[Dict] = None,
                 fix: Optional[Fixer] = None):
        self.description = description
        self.details = details or {}
        self.fix = fix

    def get_summary(self) -> str:
        return self.description

class HealthCheckResult(NamedTuple):
    severity: Severity
    issue: Issue | None = None

class HealthCheck:
    """
    Wrapper for a health check function.
    """
    def __init__(self, check_callable: Callable[[], HealthCheckResult], packages: Optional[list[str]] = None):
        """
        :param check_callable: A function that performs a health check and returns a HealthCheckResult or None (if no issue was found).
        """
        self.check_callable = check_callable
        if packages:
            cache.add(packages)

    def run(self) -> HealthCheckResult:
        """
        Execute the check callable and return its result directly.

        :return: A HealthCheckResult if an issue is found .
        """
        return self.check_callable()

class PackageCache:
    """
    A simple cache to store package existence checks.
    """
    def __init__(self):
        self.cache: Dict[str, PackageInfo] = {}
        self.pending: Set[str] = set()
        self.cached_kernels: Dict[str, Tuple[Path, str]] = {}

    def add(self, packages: list[str]):
        """
        Ask for a list of packages to be checked.

        :param packages: A list of package names to cache.
        """
        self.pending.update(packages)

    def get(self, package_name: str) -> Optional[PackageInfo]:
        """
        Get the cached PackageInfo for a package name.

        :param package_name: The name of the package to retrieve.
        :return: A PackageInfo object if found, otherwise None.
        """
        return self.cache.get(package_name)

    def execute(self) -> None:
        """
        Execute the cache, checking all pending packages.
        This should be called before any checks that require package existence.
        """
        self.cache.clear()
        if not self.pending:
            return

        packages = " ".join(self.pending)
        result = run_shell_command(f"LANG=C LANGUAGE=C LC_ALL=C pacman -Qq {packages}")

        if result.exit_code != 0 and result.exit_code != 1:
            raise Exception(f"Failed to check package existence with pacman: {result.stderr}")

        for package_name in self.pending:
            self.cache[package_name] = PackageInfo(
                installed=True,
                provided=True,
            )

        regex_pattern = re.compile(f"^error: package '([^']+)' was not found$")

        for line in result.stderr.splitlines():
            if match := regex_pattern.match(line):
                package_name = match.group(1)
                self.cache[package_name].installed = False

        for line in result.stdout.splitlines():
            name = line.strip()
            if name in self.cache:
                self.cache[name].provided = False

        # Find all kernels that are actual installed packages
        result = run_shell_command(f"LANG=C LANGUAGE=C LC_ALL=C pacman -Qo /usr/lib/modules/*/pkgbase")
        if result.exit_code != 0 and result.exit_code != 1:
            raise Exception(f"Failed to check kernel module existence with pacman: {result.stderr}")

        kernel_regex_pattern = re.compile(r"^(/usr/lib/modules/[^/]+)/pkgbase is owned by ([^ ]+) ([^ ]+)$")
        for line in result.stdout.splitlines():
            line = line.strip()
            if not line:
                continue
            if match := kernel_regex_pattern.match(line):
                kernel_path = Path(match.group(1))
                package_name = match.group(2)
                package_version = match.group(3)
                self.cached_kernels[package_name] = (kernel_path, package_version)

class Fixer:
    """A class to encapsulate a fix action, its priority, and other metadata."""

    def __init__(
        self,
        fix_function: Callable[[Issue], Tuple[bool, str]],
        priority: int = 5,
        excludes: list[str] = None,
        actions: list[str] = None,
    ):
        self.fix_function = fix_function
        self.priority = priority
        self.excludes = excludes if excludes is not None else []
        self.actions = (actions if actions is not None else [])

    def fix(self, issue: Issue) -> Tuple[bool, str]:
        return self.fix_function(issue)

class ShellResult(NamedTuple):
    exit_code: int; stdout: str; stderr: str

class SystemdStatus(NamedTuple):
    active_state: str; is_enabled: bool

# -- Convenience Functions --

cache = PackageCache()
is_virtual_machine = False

def run_shell_command(command: str, timeout: int = 10) -> ShellResult:
    """
    Run a command in a bash shell and capture its output.

    :param command: The shell command string to execute.
    :param timeout: The maximum time to wait for a command to complete.
    :return: A ShellResult object containing the exit code, stdout, and stderr.
    """
    try:
        process = subprocess.run(
            ['/bin/bash', '-c', command],
            capture_output=True, text=True, timeout=timeout, check=False
        )
        return ShellResult(process.returncode, process.stdout.strip(), process.stderr.strip())
    except FileNotFoundError:
        return ShellResult(127, "", "Shell (/bin/bash) not found.")
    except subprocess.TimeoutExpired:
        return ShellResult(124, "", f"Command timed out after {timeout} seconds.")

def get_systemd_service_status(service_name: str) -> Optional[SystemdStatus]:
    """
    Check the status and enablement state of a systemd service.

    :param service_name: The name of the service (e.g., "sshd.service").
    :return: A SystemdStatus object if the service is found, otherwise None.
    """
    exists_result = run_shell_command(f"systemctl list-unit-files {service_name}")
    if exists_result.exit_code != 0:
        return None  # Service does not exist
    active_result = run_shell_command(f"systemctl is-active {service_name}")
    enabled_result = run_shell_command(f"systemctl is-enabled {service_name}")
    return SystemdStatus(active_result.stdout, enabled_result.exit_code == 0)
      
def run_interactive_command(command: str, require_consent: bool = False) -> int:
    """
    Runs a command in a bash shell inside a pseudo-terminal (pty).

    This is the ideal way to run programs that require user interaction, have complex UIs.
    For example: pacman.

    :param command: The command string to execute in the bash shell.
    :param require_consent: If True, prompt the user for [y/N] confirmation
                            before running the command. Defaults to False.
    :return: The integer exit code of the command.
             - Returns 130 if the user cancels at the consent prompt.
             - Returns 127 if /bin/bash is not found.
             - Returns 1 for other errors during execution.
    """
    if require_consent:
        print(f"{COLORS['LOW']}The following command will be executed:{COLORS['RESET']}")
        print(f"  {command}")
        
        try:
            user_input = input("Continue? [y/N]: ")
        except (KeyboardInterrupt, EOFError):
            return 130

        if user_input.strip().lower() != 'y':
            print("Operation cancelled by user.")
            return 130

    argv = ['/bin/bash', '-c', command]
    
    try:
        wait_status = pty.spawn(argv)
    except FileNotFoundError:
        return 127
    except Exception as e:
        print(f"{COLORS['CRITICAL']}--- An error occurred: {e} ---{COLORS['RESET']}", file=sys.stderr)
        return 1
    
    if os.WIFEXITED(wait_status):
        exit_code = os.WEXITSTATUS(wait_status)
        return exit_code
    else:
        return 1

# -- Fix Functions --

def _fix_enable_services(issue: Issue) -> Tuple[bool, str]:
    """
    Generic fix to enable a systemd service.

    :param issue: The Issue object containing details about the service.
    :return: A tuple of (success_boolean, status_message).
    """
    services = issue.details.get("services")

    if not services:
        return False, "No services provided to enable."
    for service in services:
        result = run_interactive_command(f"systemctl enable {service}", require_consent=True)
        if result != 0:
            return False, f"Failed to enable service {service}. Exit code: {result}"
    return True, f"Enabled services: {', '.join(services)} successfully."

fix_enable_services = Fixer(fix_function=_fix_enable_services, priority=1)


def _fix_pacman_conf(issue: Issue) -> Tuple[bool, str]:
    """
    Fix the pacman configuration by running a remote fix command.

    :param issue: The Issue object containing details about the pacman configuration.
    :return: A tuple of (success_boolean, status_message).
    """
    result = run_interactive_command("garuda-update remote fix", require_consent=True)
    if result == 0:
        return True, "Pacman configuration fixed successfully."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to fix pacman configuration. Exit code: {result}"

fix_pacman_conf = Fixer(fix_function=_fix_pacman_conf, priority=9, excludes=["pacman_conf", "update"], actions=["pacman_conf", "update"])


def _fix_install_nvidia_drivers(issue: Issue) -> Tuple[bool, str]:
    """
    Install the nvidia drivers using garuda-hardware-tool.

    :param issue: The Issue object
    :return: A tuple of (success_boolean, status_message).
    """
    result = run_interactive_command("sudo garuda-hardware-tool --nonfree", require_consent=True)
    if result == 0:
        return True, "Nvidia drivers installed successfully. Please reboot your system."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to install nvidia drivers. Exit code: {result}"

fix_install_nvidia_drivers = Fixer(fix_function=_fix_install_nvidia_drivers, priority=7)

def _fix_install_packages(issue: Issue) -> Tuple[bool, str]:
    """
    Install missing packages using pacman.

    :param issue: The Issue object containing details about the missing packages.
    :return: A tuple of (success_boolean, status_message).
    """
    packages = issue.details.get("packages")
    if not packages:
        return False, "No packages provided to install."

    package_list = " ".join(packages)
    result = run_interactive_command(f"sudo pacman -S {package_list}", require_consent=True)
    if result == 0:
        return True, f"Installed packages: {package_list} successfully."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to install packages. Exit code: {result}"

fix_install_packages = Fixer(fix_function=_fix_install_packages, priority=7)

def _fix_update_system(issue: Issue) -> Tuple[bool, str]:
    """
    Update the system using garuda-update.

    :param issue: The Issue object
    :return: A tuple of (success_boolean, status_message).
    """
    result = run_interactive_command("garuda-update", require_consent=False)
    if result == 0:
        return True, "System updated successfully."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to update system. Exit code: {result}"

fix_update_system = Fixer(fix_function=_fix_update_system, priority=8, excludes=["update"], actions=["update"])

def _fix_remove_old_btrfs_snapshots(issue: Issue) -> Tuple[bool, str]:
    """
    Remove old Btrfs snapshots using snapper-tools.

    :param issue: The Issue object
    :return: A tuple of (success_boolean, status_message).
    """
    result = run_interactive_command("sudo snapper-tools delete-old", require_consent=True)
    if result == 0:
        return True, "Old Btrfs snapshots removed successfully."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to remove old Btrfs snapshots. Exit code: {result}"

fix_remove_old_btrfs_snapshots = Fixer(fix_function=_fix_remove_old_btrfs_snapshots, priority=1)

def _fix_time_sync(issue: Issue) -> Tuple[bool, str]:
    """
    Enable NTP synchronization using timedatectl.

    :param issue: The Issue object
    :return: A tuple of (success_boolean, status_message).
    """
    result = run_interactive_command("sudo timedatectl set-ntp true", require_consent=True)
    if result == 0:
        return True, "NTP synchronization enabled successfully."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to enable NTP synchronization. Exit code: {result}"

fix_time_sync = Fixer(fix_function=_fix_time_sync, priority=10, excludes=[], actions=[])

def _fix_dracut_rebuild(issue: Issue) -> Tuple[bool, str]:
    """
    Rebuild the initramfs using dracut.

    :param issue: The Issue object
    :return: A tuple of (success_boolean, status_message).
    """
    if os.path.exists("/usr/bin/update-grub") and os.path.exists("/boot/grub/grub.cfg"):
        result = run_interactive_command("sudo dracut-rebuild && sudo update-grub", require_consent=True)
    else:
        result = run_interactive_command("sudo dracut-rebuild", require_consent=True)
        print(f"{COLORS['INFO']}Note: It may be necessary to manually update the system's bootloader configuration to prevent boot issues.{COLORS['RESET']}")
    if result == 0:
        return True, "Initramfs rebuilt successfully using dracut."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to rebuild initramfs using dracut. Exit code: {result}"

fix_dracut_rebuild = Fixer(fix_function=_fix_dracut_rebuild, priority=1, excludes=["initramfs"], actions=["initramfs"])

def _fix_snapper_reset_config(issue: Issue) -> Tuple[bool, str]:
    """
    Reset the snapper configuration to default.

    :param issue: The Issue object
    :return: A tuple of (success_boolean, status_message).
    """
    result = run_interactive_command("garuda-update remote reset-snapper", require_consent=True)
    if result == 0:
        return True, "Snapper configuration reset to default successfully."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to reset snapper configuration. Exit code: {result}"

fix_snapper_reset_config = Fixer(fix_function=_fix_snapper_reset_config, priority=10, excludes=["snapper"], actions=["snapper"])

def _fix_dkms_rebuild(issue: Issue) -> Tuple[bool, str]:
    """
    Rebuild all DKMS modules.
    :param issue: The Issue object
    :return: A tuple of (success_boolean, status_message).
    """
    result = run_interactive_command("sudo dkms-rebuild", require_consent=True)
    if result == 0:
        return True, "DKMS modules were rebuilt"
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to rebuild DKMS modules. Exit code: {result}"
# Needs to execute before dracut to prevent depline issues
fix_dkms_rebuild = Fixer(fix_function=_fix_dkms_rebuild, priority=2, excludes=["dkms"], actions=["dkms"])

def _fix_remove_packages(issue: Issue) -> Tuple[bool, str]:
    """
    Remove deprecated or unwanted packages using pacman.

    :param issue: The Issue object containing details about the packages to remove.
    :return: A tuple of (success_boolean, status_message).
    """
    packages = issue.details.get("packages")
    if not packages:
        return False, "No packages provided to install."

    if "steam-native-runtime" in packages:
        steam_explicit = run_shell_command("pacman -Qqe steam")
        if steam_explicit.exit_code != 0:
            result = run_interactive_command(f"sudo pacman -D --asexplicit steam", require_consent=True)
            if result == 0:
                print("Set package 'steam' as explicitly installed.")
            elif result == 130:
                return False, "Fix cancelled by user."
            else:
                return False, f"Failed to remove packages. Exit code: {result}"

    package_list = " ".join(packages)
    result = run_interactive_command(f"sudo pacman -Rs {package_list}", require_consent=True)
    if result == 0:
        return True, f"Removed packages: {package_list} successfully."
    elif result == 130:
        return False, "Fix cancelled by user."
    else:
        return False, f"Failed to remove packages. Exit code: {result}"

fix_remove_packages = Fixer(fix_function=_fix_remove_packages, priority=7)

# -- Check Implementations --

def check_for_packages(severity: Severity, package_names: list[str], package_names_providable: list[str], text: str) -> Optional[HealthCheckResult]:
    """
    Check if garuda linux packages are installed.

    :param severity: The severity level for the check.
    :param package_names: A list of package names to check.
    :return: A HealthCheckResult if any package is missing .
    """
    missing_packages = [pkg for pkg in package_names if not ((pkg_info := cache.get(pkg)).installed and not pkg_info.provided)]
    missing_packages += [pkg for pkg in package_names_providable if not ((pkg_info := cache.get(pkg)).installed)]
    
    if missing_packages:
        return HealthCheckResult(
            severity,
            Issue(
                f"{text}: {', '.join(missing_packages)}",
                fix=fix_install_packages,
                details={"packages": missing_packages}
            )
        )
    return HealthCheckResult(Severity.NO_ISSUE)

def check_for_core_packages() -> HealthCheck:
    """
    Check if core garuda linux packages are installed.

    :return: A HealthCheckResult if a core package is missing .
    """
    core_packages = ["base", "chaotic-keyring", "chaotic-mirrorlist"]
    def check() -> HealthCheckResult:
        return check_for_packages(
            Severity.CRITICAL, core_packages, [], "Core garuda linux packages are missing"
        )
    return HealthCheck(check, core_packages)

def check_for_important_packages() -> HealthCheck:
    """
    Check if important garuda linux packages are installed.

    :return: A HealthCheckResult if an important package is missing .
    """
    important_packages = ["garuda-common-settings", "garuda-update", "garuda-libs", "garuda-config-agent"]
    def check() -> HealthCheckResult:
        return check_for_packages(
            Severity.HIGH, important_packages, [], "Important garuda linux packages are missing"
        )
    return HealthCheck(check, important_packages)

def check_for_optional_packages() -> HealthCheck:
    """
    Check if optional garuda linux packages are installed.

    :return: A HealthCheckResult if an optional package is missing .
    """
    optional_packages = ["garuda-rani", "garuda-hardware-tool", "garuda-system-maintenance"]
    def check() -> HealthCheckResult:
        return check_for_packages(
            Severity.INFO, [], optional_packages, "Optional garuda linux utilities are missing"
        )
    return HealthCheck(check, optional_packages)

def check_for_deprecated_packages() -> HealthCheck:
    """
    Check if any deprecated packages are installed
    """
    deprecated_packages = ["garuda-hotfixes", "steam-native-runtime"]
    def check() -> HealthCheckResult:
        found = [pkg for pkg in deprecated_packages if ((pkg_info := cache.get(pkg)).installed and not pkg_info.provided)]
        if found:
          return HealthCheckResult(
            Severity.LOW,
            Issue(
                f"Deprecated/Outdated/Removed packages should be removed: {', '.join(found)}",
                fix=fix_remove_packages,
                details={"packages": found}
            )
        )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, deprecated_packages)

def check_pacman_conf() -> HealthCheck:
    """
    Check if the pacman configuration file exists and has all the correct repositories.

    :return: A HealthCheckResult if the configuration is missing or incorrect .
    """
    def check() -> HealthCheckResult:
        pacman_conf_path = "/etc/pacman.conf"
        if not os.path.exists(pacman_conf_path):
            return HealthCheckResult(
                Severity.CRITICAL,
                Issue(
                    "Pacman configuration file is missing",
                    fix=fix_pacman_conf
                )
            )

        found_repos = set()
        # Check if the file contains the required repositories
        with open(pacman_conf_path, 'r') as f:
            for line in f:
                stripped_line = line.strip()
                if stripped_line.startswith('[') and stripped_line.endswith(']'):
                    found_repos.add(stripped_line)
        required_repos = ["[garuda]", "[core]", "[extra]", "[chaotic-aur]"]
        if not all(repo in found_repos for repo in required_repos):
            return HealthCheckResult(
                Severity.CRITICAL,
                Issue(
                    "Pacman configuration file is missing required repositories",
                    fix=fix_pacman_conf
                )
            )
        if not any(repo == "[multilib]" for repo in found_repos):
            return HealthCheckResult(
                Severity.INFO,
                Issue(
                    "Pacman configuration file does not have the multilib repository enabled",
                    fix=None
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_nvidia_drivers() -> HealthCheck:
    """
    Check if nvidia drivers are installed.

    :return: A HealthCheckResult if proprietary drivers are missing .
    """
    def check() -> HealthCheckResult:
        if not os.path.exists("/usr/bin/garuda-hardware-tool"):
            return HealthCheckResult(Severity.NOT_APPLICABLE)
        result = run_shell_command("garuda-hardware-tool --nonfree --check")
        if result.exit_code != 0:
            return HealthCheckResult(Severity.ERROR)
        if "nvidia" in result.stdout.lower():
            if cache.get("NVIDIA-MODULE").installed:
                if cache.get("GARUDA-HARDWARE-PROFILE-NVIDIA").installed:
                    return HealthCheckResult(
                        Severity.INFO,
                        Issue(
                            "Nvidia drivers are installed, but an incorrect or a suboptimal hardware profile is active",
                            fix=fix_install_nvidia_drivers
                        )
                    )
                else:
                    return HealthCheckResult(
                        Severity.INFO,
                        Issue(
                            "Nvidia drivers are installed, but no hardware profile is active",
                            fix=fix_install_nvidia_drivers
                        )
                    )
            else:
                return HealthCheckResult(
                    Severity.INFO,
                    Issue(
                        "Nvidia drivers are not installed",
                        fix=fix_install_nvidia_drivers
                    )
                )

        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, ["NVIDIA-MODULE", "GARUDA-HARDWARE-PROFILE-NVIDIA"])

def check_nvidia_kernel_module() -> HealthCheck:
    """
    Check if the nvidia kernel module is loaded.

    :return: A HealthCheckResult if the nvidia kernel module is not loaded .
    """
    def check() -> HealthCheckResult:
        if cache.get("NVIDIA-MODULE").installed:
            # Check if "nvidia" kernel module is loaded
            if not os.path.exists("/sys/module/nvidia/"):
                return HealthCheckResult(
                    Severity.HIGH,
                    Issue(
                        "Nvidia kernel module is not loaded",
                        fix=None
                    )
                )
        # This is not NOT_APPLICABLE because if no NVIDIA modules are installed, they can not possibly be malfunctioning.
        # Therefore: NO_ISSUE
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, ["NVIDIA-MODULE"])

def check_nvidia_kernel_modules() -> HealthCheck:
    """
    Check if all the installed kernels have the nvidia kernel module installed.
    Compared to check_nvidia_kernel_module, this applies only after a reboot, while
    check_nvidia_kernel_module checks the currently loaded kernel module.
    """
    def has_nvidia_dkms() -> bool:
        """
        Check if any nvidia kernel module sources exist in /usr/src
        """
        dkms_path = Path("/usr/src")
        nvidia_modules = list(dkms_path.glob("nvidia*"))
        return any(nvidia_modules)

    def check() -> HealthCheckResult:
        # This is not NOT_APPLICABLE because if no NVIDIA modules are installed, they can not possibly be malfunctioning.
        # Therefore: NO_ISSUE
        if not cache.get("NVIDIA-MODULE").installed:
            return HealthCheckResult(Severity.NO_ISSUE)

        missing_kernels = []
        for kernel, data in cache.cached_kernels.items():
            kernel_path, package_version = data
            dkms_path = kernel_path / "updates" / "dkms"
            package_path = kernel_path / "extramodules"

            dkms_modules = dkms_path.glob("nvidia.ko*")
            if any(dkms_modules):
                continue

            package_modules = package_path.glob("nvidia.ko*")
            if any(package_modules):
                continue

            main_modules = kernel_path.glob("nvidia.ko*")
            if any(main_modules):
                continue

            missing_kernels.append(f"{kernel}")

        if missing_kernels:
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "Nvidia kernel module is not installed for the following kernels: " + ", ".join(missing_kernels),
                    fix=fix_dkms_rebuild if has_nvidia_dkms() else None,
                )
            )

        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, ["NVIDIA-MODULE"])

def check_btrfs_root_storage() -> HealthCheck:
    """
    Check if the root filesystem is using Btrfs and has enough free space.

    :return: A HealthCheckResult if the root filesystem is not Btrfs .
    """
    def check() -> HealthCheckResult:
        if (not os.path.exists("/usr/bin/btrfs")):
            return HealthCheckResult(Severity.NOT_APPLICABLE)

        # Check free space
        free_space_result = run_shell_command("btrfs filesystem usage / -m --si")
        if free_space_result.exit_code != 0:
            return HealthCheckResult(Severity.ERROR)
        # Find Overall -> Device unallocated:
        match = re.search(r"\s+Free \(estimated\):\s+(\d+)\.\d+MB", free_space_result.stdout)
        if match:
            free_space_mb = int(match.group(1))
            if free_space_mb < 3000:  # Less than 5GB free space
                return HealthCheckResult(
                    Severity.HIGH,
                    Issue(
                        "Root Btrfs filesystem has less than 5GB free space",
                        fix=None
                    )
                )
            elif free_space_mb < 10000:  # Less than 10GB free space
                return HealthCheckResult(
                    Severity.LOW,
                    Issue(
                        "Root Btrfs filesystem has less than 10GB free space",
                        fix=None
                    )
                )

        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_btrfs_timers() -> HealthCheck:
    """
    Check if the Btrfs maintenance timers are enabled.

    :return: A HealthCheckResult if the btrfs maintenance services are not running or enabled .
    """
    def check() -> HealthCheckResult:
        services = ["btrfs-scrub.timer", "btrfs-balance.timer", "btrfs-defrag.timer"]
        missing_services = []
        for service in services:
            status = get_systemd_service_status(service)
            if not status:
                return HealthCheckResult(Severity.ERROR)
            if status.is_enabled is False:
                missing_services.append(service)

        if missing_services:
            issue_description = "Btrfs maintenance timers are not enabled: " + ", ".join(missing_services)
            return HealthCheckResult(
                Severity.LOW,
                Issue(
                    issue_description,
                    details={"services": missing_services},
                    fix=fix_enable_services
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_using_dracut() -> HealthCheck:
    """
    Check if the system is using dracut for initramfs.

    :return: A HealthCheckResult if dracut is not being used .
    """
    def check() -> HealthCheckResult:
        if not cache.get("dracut").installed:
            return HealthCheckResult(
                Severity.INFO,
                Issue(
                    "System is not using dracut for initramfs",
                    fix=fix_install_packages,
                    details={"packages": ["garuda-dracut-support"]}
                )
            )
        if not ((pkg := cache.get("garuda-dracut-support")).installed and not pkg.provided):
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "System is using dracut, but garuda-dracut-support is not installed",
                    fix=fix_install_packages,
                    details={"packages": ["garuda-dracut-support"]}
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, ["dracut", "garuda-dracut-support"])

def check_using_pipewire() -> HealthCheck:
    """
    Check if the system is using PipeWire for audio.

    :return: A HealthCheckResult if PipeWire is not being used .
    """
    def check() -> HealthCheckResult:
        if not cache.get("pipewire").installed:
            return HealthCheckResult(
                Severity.INFO,
                Issue(
                    "System is not using PipeWire for audio",
                    fix=fix_install_packages,
                    details={"packages": ["pipewire-support"]}
                )
            )

        if not ((pkg := cache.get("pipewire-support")).installed and not pkg.provided):
            return HealthCheckResult(
                Severity.LOW,
                Issue(
                    "System is using PipeWire, but pipewire-support is not installed",
                    fix=fix_install_packages,
                    details={"packages": ["pipewire-support"]}
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, ["pipewire", "pipewire-support"])

def check_partial_upgrade() -> HealthCheck:
    """
    Check if the system is partially upgraded.

    :return: A HealthCheckResult if the system is partially upgraded .
    """
    def check() -> HealthCheckResult:
        partial_upgrade_file = "/var/lib/garuda/partial_upgrade"

        if os.path.exists(partial_upgrade_file):
            return HealthCheckResult(
                Severity.CRITICAL,
                Issue(
                    "System is partially upgraded",
                    fix=fix_update_system
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_outdated_system() -> HealthCheck:
    """
    Check if the system is outdated. (older than 2 weeks)

    :return: A HealthCheckResult if the system is outdated .
    """
    def check() -> HealthCheckResult:
        last_update_file = "/var/lib/garuda/last_update"
        # File timestamp is the date of the last update
        if not os.path.exists(last_update_file):
            return HealthCheckResult(Severity.NOT_APPLICABLE)
        last_update_time = os.path.getmtime(last_update_file)
        current_time = time.time()
        # Check if the last update was more than 2 weeks ago
        if (current_time - last_update_time) > (14 * 24 * 60 * 60):  # 2 weeks in seconds
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "System is outdated (last update more than 2 weeks ago)",
                    fix=fix_update_system
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_pending_reboot() -> HealthCheck:
    """
    Check if a reboot is pending.

    :return: A HealthCheckResult if a reboot is pending .
    """
    def check() -> HealthCheckResult:
        # Check if last update is newer than the last reboot
        last_update_file = "/var/lib/garuda/last_update"
        last_reboot_benchmark = "/proc"

        if not os.path.exists(last_update_file):
            return HealthCheckResult(Severity.NOT_APPLICABLE)

        last_update_time = os.path.getmtime(last_update_file)
        last_reboot_time = os.path.getctime(last_reboot_benchmark)

        if last_update_time > last_reboot_time:
            return HealthCheckResult(
                Severity.INFO,
                Issue(
                    "A reboot is pending (update applied since last reboot)",
                    fix=None
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_failed_systemd_services() -> HealthCheck:
    """
    Check for failed systemd services.

    :return: A HealthCheckResult if any failed services are found .
    """
    def check() -> HealthCheckResult:
        result = run_shell_command("systemctl list-units --failed --full --all --plain --no-legend")
        if result.exit_code != 0 or not result.stdout.strip():
            # No failed services
            return HealthCheckResult(Severity.NO_ISSUE)

        failed_services = []
        for line in result.stdout.strip().splitlines():
            failed_services.append(line.split()[0])

        if failed_services:
            return HealthCheckResult(
                Severity.LOW,
                Issue(
                    "Failed systemd services found: " + ", ".join(failed_services),
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_has_old_btrfs_snapshots() -> HealthCheck:
    """
    Check if there are old Btrfs snapshots that can be deleted.

    :return: A HealthCheckResult if old snapshots are found .
    """
    def check() -> HealthCheckResult:
        if not os.path.exists("/usr/bin/snapper-tools"):
            return HealthCheckResult(Severity.NOT_APPLICABLE)

        result = run_shell_command("pkexec snapper-tools find-old")
        if result.exit_code == 0:
            return HealthCheckResult(
                Severity.LOW,
                Issue(
                    "Old Btrfs snapshots found that can be deleted",
                    fix=fix_remove_old_btrfs_snapshots
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_cpu_temperature() -> HealthCheck:
    """
    Check if the CPU temperature is within safe limits.

    :return: A HealthCheckResult if the CPU temperature is too high .
    """
    def check() -> HealthCheckResult:
        thermal_zone_path = "/sys/class/thermal/"
        cpu_temp = "x86_pkg_temp"
        try:
            thermal_zones = Path(thermal_zone_path).glob("thermal_zone*")
            if not thermal_zones:
                return HealthCheckResult(Severity.NO_ISSUE)

            for zone in thermal_zones:
                type_file = zone / "type"
                if not type_file.exists():
                    continue
                with open(type_file, 'r') as f:
                    if cpu_temp != f.read().strip():
                        continue
                temp_file = zone / "temp"
                if not temp_file.exists():
                    break
                with open(temp_file, 'r') as f:
                    temp = int(f.read().strip())
                    if temp > 90000:  # Temperature in millidegrees Celsius
                        return HealthCheckResult(
                            Severity.LOW,
                            Issue(
                                f"CPU temperature is too high: {temp / 1000:.1f}°C",
                                fix=None
                            )
                        )
                    else:
                        break
        except (IOError, ValueError):
            return HealthCheckResult(Severity.ERROR)
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_time_sync() -> HealthCheck:
    """
    Check if the system time is synchronized.

    :return: A HealthCheckResult if the system time is not synchronized .
    """
    def check() -> HealthCheckResult:
        result = run_shell_command("timedatectl show --property=NTP --value")
        if result.exit_code != 0 or not result.stdout.strip():
            return HealthCheckResult(Severity.ERROR)
        if result.stdout.strip() != "yes":
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "NTP is not enabled, system time may not be synchronized",
                    fix=fix_time_sync
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_smart_status() -> HealthCheck:
    """
    Check the SMART status of all disks.

    :return: A HealthCheckResult if any disk has a failing SMART status .
    """
    def check() -> HealthCheckResult:
        if (not os.path.exists("/usr/bin/smartctl") or os.getuid() != 0 or is_virtual_machine):
            return HealthCheckResult(Severity.NOT_APPLICABLE)
        result = run_shell_command("smartctl --scan")
        if result.exit_code != 0:
            return HealthCheckResult(Severity.ERROR)
        disks = result.stdout.strip().splitlines()
        if not disks:
            return HealthCheckResult(Severity.NOT_APPLICABLE)

        for disk in disks:
            device = disk.split()[0]
            smart_result = run_shell_command(f"smartctl -H {device}")
            expected_fail = "SMART overall-health self-assessment test result: FAILED"
            if expected_fail in smart_result.stdout or expected_fail in smart_result.stderr:
                return HealthCheckResult(
                    Severity.CRITICAL,
                    Issue(
                        f"SMART status for {device} is failing",
                        fix=None,
                    )
                )
            if smart_result.exit_code != 0:
                continue
            if "SMART overall-health self-assessment test result: PASSED" not in smart_result.stdout:
                return HealthCheckResult(
                    Severity.CRITICAL,
                    Issue(
                        f"SMART status for {device} is failing",
                        fix=None,
                    )
                )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_available_ram() -> HealthCheck:
    """
    Check if the system has enough available RAM.

    :return: A HealthCheckResult if the available RAM is below a threshold .
    """
    def check() -> HealthCheckResult:
        meminfo_path = "/proc/meminfo"
        try:
            data = {}
            with open(meminfo_path, 'r') as f:
                pattern = re.compile(r"^([^:]+):\s+(\d+)\s+kB$")
                for line in f:
                    match = pattern.match(line)
                    if match:
                        key, value = match.groups()
                        data[key] = int(value)
            if "MemAvailable" not in data or "MemTotal" not in data:
                return HealthCheckResult(Severity.ERROR)
            if "SwapTotal" in data and "SwapFree" in data and data["SwapTotal"] > 0:
                if data["MemAvailable"] < data["MemTotal"] * 0.05:
                    return HealthCheckResult(
                        Severity.HIGH,
                        Issue(
                            f"Available RAM is below 5% of total RAM: {data['MemAvailable'] / 1024:.2f} MB",
                            fix=None
                        )
                    )
                # Less than 15% of total swap available AND less than 15% of total RAM available
                if (data["SwapFree"] < data["SwapTotal"] * 0.15
                        and data["MemAvailable"] < data["MemTotal"] * 0.15):
                    return HealthCheckResult(
                        Severity.HIGH,
                        Issue(
                            f"Available RAM and swap are both below 15% of total: "
                            f"{data['MemAvailable'] / 1024:.2f} MB available RAM, "
                            f"{data['SwapFree'] / 1024:.2f} MB available swap",
                            fix=None
                        )
                    )
            # No swap
            else:
                if data["MemAvailable"] < data["MemTotal"] * 0.10:
                    return HealthCheckResult(
                        Severity.HIGH,
                        Issue(
                            f"Available RAM is below 10% of total RAM: {data['MemAvailable'] / 1024:.2f} MB",
                            fix=None
                        )
                    )
        except (IOError, ValueError):
            return HealthCheckResult(Severity.ERROR)
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_initramfs_success() -> HealthCheck:
    """
    Check if the initramfs error file (/var/lib/garuda/initramfs_error) exists.

    :return: A HealthCheckResult if the initramfs build failed .
    """
    def check() -> HealthCheckResult:
        initramfs_error_file = "/var/lib/garuda/initramfs_error"
        if os.path.exists(initramfs_error_file):
            return HealthCheckResult(
                Severity.CRITICAL,
                Issue(
                    "Pacman hooks report that the initramfs build has failed. System may not boot.",
                    fix=fix_dracut_rebuild if os.path.exists("/usr/bin/dracut-rebuild") else None,
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)

    return HealthCheck(check, [])

def check_snapper_services() -> HealthCheck:
    """
    Check if the snapper services are running.

    :return: A HealthCheckResult if the snapper services are not running .
    """
    def check() -> HealthCheckResult:
        if not cache.get("snapper-support").installed:
            return HealthCheckResult(Severity.INFO, Issue(
                "Snapper (automatic btrfs system restore) is not installed and configured" if not cache.get("snapper").installed else
                "Snapper (automatic btrfs system restore) is not using the official Garuda Linux configuration",
                fix=fix_install_packages,
                details={"packages": ["snapper-support"]}
            ))
        if not (service := get_systemd_service_status("snapper-cleanup.timer")):
            return HealthCheckResult(Severity.ERROR)
        if not service.is_enabled:
            return HealthCheckResult(
                Severity.LOW,
                Issue(
                    "Snapper cleanup timer is not enabled",
                    fix=fix_enable_services,
                    details={"services": ["snapper-cleanup.timer"]}
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, ["snapper-support", "snapper"])

def check_snapper_config() -> HealthCheck:
    """
    Check if the snapper configuration is correct.

    :return: A HealthCheckResult if the snapper configuration is incorrect .
    """
    def check() -> HealthCheckResult:
        if not cache.get("snapper-support").installed:
            return HealthCheckResult(Severity.NOT_APPLICABLE)
        if not os.path.exists("/etc/snapper/configs/root"):
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "Snapper configuration for root is missing",
                    fix=fix_snapper_reset_config
                )
            )
        if not os.path.exists("/.snapshots"):
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "Snapper snapshots directory is missing",
                    fix=fix_snapper_reset_config
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, ["snapper-support"])

def check_kernels_in_boot() -> HealthCheck:
    """
    Check if all the kernels present in /usr/lib/modules are properly installed in /boot.
    """
    def check() -> HealthCheckResult:
        boot_path = Path("/boot")

        any_kernel_installed = False
        any_kernel_valid = False
        missing_kernels = []
        invalid_kernels = []

        secureboot_enabled = os.path.exists("/etc/garuda/secureboot/enabled")

        for kernel, data in cache.cached_kernels.items():
            kernel_boot_path = boot_path / ("vmlinuz-" + kernel)
            initramfs_boot_path = boot_path / ("initramfs-" + kernel + ".img")
            kernel_modules_dir, kernel_version = data
            if kernel_boot_path.exists() and initramfs_boot_path.exists():
                any_kernel_installed = True
                if secureboot_enabled or filecmp.cmp(kernel_boot_path, kernel_modules_dir / "vmlinuz", shallow=False):
                    any_kernel_valid = True
                else:
                    invalid_kernels.append(kernel)
            else:
                missing_kernels.append(kernel)

        if not any_kernel_installed:
            return HealthCheckResult(
                Severity.CRITICAL,
                Issue(
                    "No (intended) kernels are installed in /boot",
                    fix=fix_dracut_rebuild if os.path.exists("/usr/bin/dracut-rebuild") else None,
                )
            )

        if not any_kernel_valid:
            return HealthCheckResult(
                Severity.CRITICAL,
                Issue(
                    "Kernels in /boot are invalid (do not match expected files)",
                    fix=fix_dracut_rebuild if os.path.exists("/usr/bin/dracut-rebuild") else None,
                )
            )

        if missing_kernels:
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "The following kernels are missing in /boot: " + ", ".join(missing_kernels),
                    fix=fix_dracut_rebuild if os.path.exists("/usr/bin/dracut-rebuild") else None,
                )
            )
        if invalid_kernels:
            return HealthCheckResult(
                Severity.HIGH,
                Issue(
                    "The following kernels in /boot are invalid: " + ", ".join(invalid_kernels),
                    fix=fix_dracut_rebuild if os.path.exists("/usr/bin/dracut-rebuild") else None,
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_efi_bootmgr() -> HealthCheck:
    """
    Check if the system is using EFI and if the correct efi bootloader is in use.
    """
    def check() -> HealthCheckResult:
        if os.path.exists("/etc/garuda/garuda-health/bypass-wrong-bootloader-check"):
            return HealthCheckResult(Severity.NOT_APPLICABLE)
        if not os.path.exists("/sys/firmware/efi") or not os.path.exists("/usr/bin/efibootmgr"):
            return HealthCheckResult(Severity.NOT_APPLICABLE)
        result = run_shell_command("efibootmgr | awk 'match($0,/^BootCurrent: ([0-9a-fA-F]{4})$/,out) { current=out[1] } match($0,/^Boot([0-9a-fA-F]{4})\\*? ([^:]+)\\t/,out) { if (out[1]==current) print out[2] }'")
        if result.exit_code != 0 or not result.stdout.strip():
            return HealthCheckResult(Severity.ERROR)
        if result.stdout.strip().lower() != "garuda":
            return HealthCheckResult(
                Severity.LOW,
                Issue(
                    f'"Garuda" is not the current bootloader: https://wiki.garudalinux.org/why-garuda-bootloader',
                    fix=None
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, [])

def check_for_packagekit() -> HealthCheck:
    """
    Check if PackageKit is installed, and if it is, warn the user and suggest its removal.
    """
    packages = ["packagekit", "packagekit-qt6", "packagekit-qt5", "gnome-packagekit"]
    def check() -> HealthCheckResult:
        pending_removals = []
        for pkg_name in packages:
            pkg = cache.get(pkg_name)
            if pkg.installed and not pkg.provided:
                pending_removals.append(pkg_name)
        if pending_removals:
            return HealthCheckResult(
                Severity.LOW,
                Issue(
                    "PackageKit is currently installed. PackageKit (Discover, GNOME Software) is not a supported method to manage packages on Garuda Linux. https://wiki.garudalinux.org/avoid-packagekit",
                    fix=fix_remove_packages,
                    details={"packages": pending_removals}
                )
            )
        return HealthCheckResult(Severity.NO_ISSUE)
    return HealthCheck(check, packages)

checks: list[HealthCheck] = [
    check_for_core_packages(),
    check_for_important_packages(),
    check_for_optional_packages(),
    check_for_deprecated_packages(),
    check_pacman_conf(),
    check_nvidia_drivers(),
    check_nvidia_kernel_module(),
    check_nvidia_kernel_modules(),
    check_btrfs_root_storage(),
    check_btrfs_timers(),
    check_using_dracut(),
    check_using_pipewire(),
    check_partial_upgrade(),
    check_outdated_system(),
    check_pending_reboot(),
    check_failed_systemd_services(),
    check_has_old_btrfs_snapshots(),
    # Unreliable/false positives on some systems
    # check_cpu_temperature(),
    check_time_sync(),
    check_smart_status(),
    check_available_ram(),
    check_initramfs_success(),
    check_snapper_services(),
    check_snapper_config(),
    check_kernels_in_boot(),
    check_efi_bootmgr(),
    check_for_packagekit(),
]

# -- Main Execution Logic --
CheckResults: TypeAlias = Dict[Severity, list[Issue]]

def perform_check_run(debug: bool) -> Tuple[CheckResults, float]:
    """
    Run all checks and collect results.

    :return: A tuple containing (results_dictionary, time_taken).
    """
    # Refresh the package cache
    time_full_start = time.time()
    cache.execute()
    results: CheckResults = {s: [] for s in Severity}
    for check in checks:
        if (debug):
            time_start = time.time()
        try:
            result = check.run()
            if (debug):
                time_end = time.time()
                print(f"Check took {time_end - time_start:.2f} seconds")
            if not result:
                raise ValueError(f"Check returned None or empty result.")
            if debug and not result.severity in printable_severities:
                print(f"Check returned a result with severity {result.severity.name}, which is not printable.")
            results[result.severity].append(result.issue)
        except Exception as e:
            results[Severity.ERROR].append(None)
            if debug:
                print(f"Error running check: {e}")
    time_full = time.time() - time_full_start
    return results, time_full

def parse_results(results: CheckResults) -> Tuple[int, int]:
    """
    Parse the results and count printable + run checks

    :param results: The dictionary of results from perform_check_run.
    :return: A tuple containing (printable_issues, run_checks, error_level_issues).
    """
    printable_issues = sum(len(results[s]) for s in printable_severities)
    info_level_issue = len(results[Severity.INFO])
    run_checks = printable_issues + len(results[Severity.NO_ISSUE])
    error_level_issues = printable_issues - info_level_issue
    return printable_issues, run_checks, error_level_issues

def print_text_report(printable_issues: int, error_level_issues: int, run_checks: int, results: CheckResults, time) -> None:
    """
    Print a text report to the console based on check results.

    :param has_issues: Boolean indicating if any issues were found.
    :param results: The dictionary of results from perform_check_run.
    """
    print(f"{COLORS['INFO']}--- System Health Check Report ---{COLORS['RESET']}")
    print(f"{COLORS['GREEN']}{run_checks}/{len(checks)} checks run in {time:.2f} seconds ⌛{COLORS['RESET']}")
    print(f"{COLORS['GREEN']}Powered by garuda-health 🦅{COLORS['RESET']}")
    if run_checks == 0:
        print(f"{COLORS['CRITICAL']}All checks failed with unknown errors.{COLORS['RESET']}")
        return
    if printable_issues == 0:
        print(f"\n✅ System health check passed. No issues found.")
        return

    has_fix=False

    for severity in reversed(printable_severities):
        if results[severity]:
            color, reset = COLORS.get(severity.name, ""), COLORS["RESET"]
            print(f"\n{color}--- {severity.name} ---{reset}")
            for res in results[severity]:
                fix_available = ""
                if res.fix:
                    fix_available = " (fix available)"
                    has_fix = True
                print(f" - {res.get_summary()}{fix_available}")

    if error_level_issues == 0:
        print(f"\n✅ System health check passed. No issues found.")
    if has_fix:
        if (error_level_issues != 0):
            print()
        print(f"{COLORS['INFO']}Run garuda-health --fix to apply fixes.{COLORS['RESET']}")

def apply_fixes(results: CheckResults) -> bool:
    """
    Attempt to apply all available fixes for the found issues.

    Fixes are applied in descending order of priority. If a fix succeeds, its
    `excludes` are noted. Subsequent fixes will be skipped if their `actions`
    overlap with any previously noted exclusions.

    :param results: The dictionary of results from perform_check_run.
    :return: True if all fixes ran without exceptions, False otherwise.
    """
    print("\n--- Applying Fixes ---")
    had_errors = False
    excluded_actions: Set[str] = set()

    fixable_results = [
        res for sev_list in results.values()
        for res in sev_list if res is not None and res.fix is not None
    ]

    fixable_results.sort(key=lambda r: r.fix.priority, reverse=True)

    for result in fixable_results:
        if not excluded_actions.isdisjoint(result.fix.actions):
            continue

        print(f'Fixing "{result.get_summary()}": ', flush=True)
        try:
            success, message = result.fix.fix(result)
            print(message)

            if success:
                if result.fix.excludes:
                    excluded_actions.update(result.fix.excludes)

        except Exception as e:
            print(f"Failed with exception: {e}")
            had_errors = True

    if not fixable_results:
        print("No issues with available fixes were found.")
        return False

    return not had_errors

def main():
    """
    Parse arguments, run all checks, and report the results.
    """
    parser = argparse.ArgumentParser(description="A flexible Linux system health checking tool.")
    parser.add_argument('--fix', action='store_true', help="Attempt to fix issues and re-run checks.")
    parser.add_argument('--json', action='store_true', help="Output in JSON. Incompatible with --fix.")
    parser.add_argument('--debug', action='store_true', help="Enable debug output.")
    parser.add_argument('--forcetty', action='store_true', help="Force color output even when not connected to a TTY.")
    args = parser.parse_args()

    if not sys.stdout.isatty() and not args.forcetty:
        for key in COLORS: COLORS[key] = ""

    if args.json and args.fix:
        parser.error("--json and --fix are incompatible.")

    ## Demand root when --fix is used
    #if args.fix and os.geteuid() != 0:
    #    print(f"{COLORS['CRITICAL']}Error: --fix requires root privileges.{COLORS['RESET']}")
    #    sys.exit(1)

    detect_virt = run_shell_command("systemd-detect-virt --quiet")
    global is_virtual_machine
    is_virtual_machine = (detect_virt.exit_code == 0)

    initial_results, initial_time = perform_check_run(args.debug)
    printable_issues, run_checks, error_level_issues = parse_results(initial_results)

    if args.json:
        json_output = {}
        if printable_issues > 0:
            for severity in reversed(printable_severities):
                if initial_results[severity]:
                    json_output[severity.name] = [
                        {
                            "description": res.get_summary(),
                            "fix_available": res.fix is not None,
                        } for res in initial_results[severity]
                    ]
        print(json.dumps(json_output, indent=2))
        sys.exit(1 if error_level_issues > 0 else 0)

    print_text_report(printable_issues, error_level_issues, run_checks, initial_results, initial_time)
    
    final_exit_code = 1 if error_level_issues > 0 else 0

    if args.fix and printable_issues > 0:
        fixes_ran_cleanly = apply_fixes(initial_results)
        
        if fixes_ran_cleanly:
            print("\n--- Re-running checks after applying fixes ---")
            final_results, final_time = perform_check_run(args.debug)
            printable_issues, run_checks, error_level_issues = parse_results(initial_results)
            
            print_text_report(printable_issues, error_level_issues, run_checks, final_results, final_time)
            
            final_exit_code = 1 if error_level_issues > 0 else 0

    sys.exit(final_exit_code)

if __name__ == "__main__":
    main()
