dotfiles/bin/helpers/functions.py

#!/usr/bin/env python3

import sys
import subprocess
import math
import random
import shutil
import datetime

try:
    import pyfiglet
except ImportError:
    pyfiglet = None

# Color codes for terminal output
COLORS = {
    "black": "\033[0;30m",
    "red": "\033[0;31m",
    "green": "\033[0;32m",
    "yellow": "\033[0;33m",
    "blue": "\033[0;34m",
    "purple": "\033[0;35m",
    "cyan": "\033[0;36m",
    "white": "\033[0;37m",
    "grey": "\033[0;90m",  # Added grey color for timestamp
    "reset": "\033[0m",
}


def printfe(color, message, show_time=True):
    """
    Print a formatted message with the specified color
    With timestamp and message type prefix similar to setup.sh
    """
    color_code = COLORS.get(color.lower(), COLORS["reset"])

    if show_time:
        # Add timestamp
        timestamp = datetime.datetime.now().strftime("%H:%M:%S")
        print(f"{COLORS['grey']}{timestamp}{COLORS['reset']}", end="")

        # Add message type based on color
        if color.lower() in ["green", "cyan", "blue", "purple"]:
            print(f"{COLORS['green']} INF {COLORS['reset']}", end="")
        elif color.lower() == "yellow":
            print(f"{COLORS['yellow']} WRN {COLORS['reset']}", end="")
        elif color.lower() == "red":
            print(f"{COLORS['red']} ERR {COLORS['reset']}", end="")

    # Print the actual message with color
    print(f"{color_code}{message}{COLORS['reset']}")


def println(message, color=None):
    """Print a line with optional color"""
    if color:
        printfe(color, message)
    else:
        printfe("reset", message)


def _rainbow_color(text, freq=0.1, offset=0):
    """Apply rainbow colors to text similar to lolcat"""
    colored_text = ""
    for i, char in enumerate(text):
        if char.strip():  # Only color non-whitespace characters
            # Calculate RGB values using sine waves with phase shifts
            r = int(127 * math.sin(freq * i + offset + 0) + 128)
            g = int(127 * math.sin(freq * i + offset + 2 * math.pi / 3) + 128)
            b = int(127 * math.sin(freq * i + offset + 4 * math.pi / 3) + 128)

            # Apply the RGB color to the character
            colored_text += f"\033[38;2;{r};{g};{b}m{char}\033[0m"
        else:
            colored_text += char

    return colored_text


def logo(continue_after=False):
    """Display the dotfiles logo"""
    try:
        # Try to read logo file first for backward compatibility
        if pyfiglet:
            # Generate ASCII art with pyfiglet and rainbow colors
            ascii_art = pyfiglet.figlet_format("Menno's Dotfiles", font="slant")
            print("\n")  # Add some space before the logo

            # Use a random offset to vary the rainbow start position
            random_offset = random.random() * 2 * math.pi
            line_offset = 0

            for line in ascii_art.splitlines():
                # Add a little variation to each line
                print(_rainbow_color(line, offset=random_offset + line_offset))
                line_offset += 0.1
        else:
            # Fallback if pyfiglet is not available
            printfe("yellow", "\n  *** Menno's Dotfiles ***\n")
            printfe("cyan", "  Note: Install pyfiglet for better logo display")
            printfe("cyan", "  (pip install pyfiglet)\n")

        if not continue_after:
            sys.exit(0)
    except Exception as e:
        printfe("red", f"Error displaying logo: {e}")


def run_command(command, shell=False):
    """Run a shell command and return the result"""
    try:
        if not shell and not shutil.which(command[0]):
            return False, f"Command '{command[0]}' not found"

        result = subprocess.run(
            command,
            shell=shell,
            check=True,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        return True, result.stdout.strip()
    except subprocess.CalledProcessError as e:
        return False, e.stderr.strip()
    except FileNotFoundError:
        return False, f"Command '{command[0]}' not found"


def command_exists(command):
    """Check if a command exists in the PATH"""
    return shutil.which(command) is not None


def ensure_dependencies():
    """Check and install required dependencies for the dotfiles system"""
    required_packages = [
        "pyfiglet",  # For ASCII art generation
    ]

    # Check if pip is available
    success, _ = run_command(["pip", "--version"])
    if not success:
        printfe(
            "red",
            "Pip is required to install missing dependencies, retry after running `dotf update`",
        )
        return False

    missing_packages = []
    for package in required_packages:
        try:
            __import__(package)
        except ImportError:
            missing_packages.append(package)

    if missing_packages:
        printfe("yellow", f"Missing dependencies: {', '.join(missing_packages)}")
        install = input("Would you like to install them now? (y/n): ").lower()
        if install == "y" or install == "yes":
            printfe("cyan", "Installing missing dependencies...")
            for package in missing_packages:
                printfe("blue", f"Installing {package}...")
                success, output = run_command(
                    ["pip", "install", "--user", package, "--break-system-packages"]
                )
                if success:
                    printfe("green", f"Successfully installed {package}")
                else:
                    printfe("red", f"Failed to install {package}: {output}")

            printfe("green", "All dependencies have been processed")
            return True
        else:
            printfe("yellow", "Skipping dependency installation")
            return False
    return True