#!/usr/bin/env python3
"""
Save Grub and Plymouth themes

Author: Arjen Balfoort
Date: 2026-02-22
"""

import sys
import os
import re
import shutil
import subprocess

# If set to True, no changes will be made
DEBUG = False


def run_command(command: list):
    """
    Run a command and print its output.
    The print is flushed to ensure that the output is immediately
    visible and can be used by the PackageManagerThread class.
    """
    print(f"Running command: {' '.join(command)}", flush=True)
    command = ["stdbuf", "-oL"] + command
    try:
        with subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1,
        ) as process:
            for line in process.stdout:
                print(line.strip(), flush=True)
            process.wait()
            if process.returncode != 0:
                print(f"Error: Command failed with return code {process.returncode}", flush=True)
                sys.exit(1)
    except Exception as e:
        print(f"Error: {e}", flush=True)
        sys.exit(1)

def is_encrypted(target: str) -> bool:
    """
    Check if a partition or mount point is LUKS-encrypted using cryptsetup.

    Args:
        target: Path to the partition (e.g., "/dev/sda1") or mount point (e.g., "/boot").

    Returns:
        bool: True if the partition is LUKS-encrypted, False otherwise.
    """
    # If target is a mount point, resolve to block device
    if os.path.isdir(target):
        try:
            result = subprocess.run(
                ['findmnt', '-n', '-o', 'SOURCE', '--target', target],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            old_target = target
            target = result.stdout.strip()
            print(f"Resolved mount point {old_target} to block device {target}")

            # Check if the device is a mapped device
            if target.startswith("/dev/mapper/"):
                return True

        except subprocess.CalledProcessError:
            return False

    # Check if the block device is LUKS-encrypted
    try:
        subprocess.run(
            ["cryptsetup", "isLuks", target],
            capture_output=True,
            check=True
        )
        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False

def clean_modules_file(modules_path: str) -> None:
    """
    Clean up the modules file by:
    - Removing leading/trailing whitespace from each line
    - Removing lines matching specific patterns (KMS, intel_agp, drm, etc.)
    """
    # Create empty modules file
    if not os.path.isfile(modules_path):
        try:
            with open(file=modules_path, mode='w', encoding='utf-8'):
                pass
            print(f"Empty file created: {modules_path}")
        except IOError:
            print(f"Failed to create empty file: {modules_path}")
            return

    patterns_to_remove = [
        r'^.*KMS$',
        r'^intel_agp$',
        r'^drm$',
        r'^nouveau modeset.*',
        r'^radeon modeset.*',
        r'^i915 modeset.*',
        r'^uvesafb\s*mode_option.*',
    ]

    try:
        with open(file=modules_path, mode='r', encoding='utf-8') as f:
            lines = f.readlines()
    except IOError:
        print(f"Error reading {modules_path}")
        return

    cleaned_lines = []
    for line in lines:
        stripped = line.strip()
        if not any(re.match(pattern, stripped) for pattern in patterns_to_remove):
            cleaned_lines.append(stripped)

    try:
        if not DEBUG:
            with open(file=modules_path, mode='w', encoding='utf-8') as f:
                f.write('\n'.join(cleaned_lines) + '\n')
        print(f"Cleaned up {modules_path}")
    except IOError:
        print(f"Failed to clean up {modules_path}")

def add_splash_to_grub(grub_default: str, enable_splash: bool) -> None:
    """
    Add or remove splash from Grub configuration.

    Args:
        grub_default (str): The path to the Grub configuration file.
        enable_splash (bool): True to add splash, False to add nosplash.
    """
    # Read the contents of the file
    with open(file=grub_default, mode='r', encoding='utf-8') as f:
        lines = f.readlines()

    # Find the line with GRUB_CMDLINE_LINUX_DEFAULT
    for i, line in enumerate(lines):
        # Extract the value inside quotes (handles both single and double quotes)
        match = re.match(r'GRUB_CMDLINE_LINUX_DEFAULT=([\'"])(.*?)\1', line)
        if not match:
            continue

        quote = match.group(1)
        value = match.group(2)

        print(f"Found GRUB_CMDLINE_LINUX_DEFAULT: {quote}{value}{quote}")

        # Remove 'splash' and 'nosplash' (as whole words)
        value = re.sub(r'\bsplash\b', '', value)
        value = re.sub(r'\bnosplash\b', '', value)

        # Add 'splash' or 'nosplash' as requested
        if enable_splash:
            value = value.strip() + ' splash'
        else:
            value = value.strip() + ' nosplash'

        print(f"Modified GRUB_CMDLINE_LINUX_DEFAULT: {value}")

        # Reconstruct the line
        lines[i] = f'GRUB_CMDLINE_LINUX_DEFAULT={quote}{value.strip()}{quote}\n'
        break

    # Write the modified lines back to the file
    try:
        if not DEBUG:
            with open(file=grub_default, mode='w', encoding='utf-8') as f:
                f.writelines(lines)
    except IOError:
        print(f"Failed to write to {grub_default}")

def plymouth_update():
    """
    Update Grub and Initramfs for Plymouth themes.
    """
    # Run update-grub in a separate thread to prevent blocking
    run_command(['update-grub'])

    # Run update-initramfs
    run_command(['update-initramfs', '-u', '-k', 'all'])

def replace_pattern_in_grub(grub_default: str,
                            pattern: str,
                            replacement: str,
                            append_if_not_exists: bool=True) -> None:
    """
    Replace all occurrences of a regex pattern with a replacement string in a file.

    Args:
        pattern (str): Regular expression pattern to search for.
        replacement (str): String to replace matches with.
    """
    try:
        with open(file=grub_default, mode='r', encoding='utf-8') as f:
            content = f.read()
    except FileNotFoundError:
        print(f"Grub configuration file not found: {grub_default}")
        return

    print(f" pattern: {pattern}, replacement: {replacement}")

    new_content, num_replacements = re.subn(pattern, replacement, content, flags=re.MULTILINE)

    if num_replacements > 0:
        try:
            if not DEBUG:
                with open(file=grub_default, mode='w', encoding='utf-8') as f:
                    f.write(new_content)
            print(f"Replaced {num_replacements} occurrences of {pattern} in {grub_default}")
        except IOError:
            print(f"Failed to replace {pattern} occurrences in {grub_default}")
    else:
        if append_if_not_exists:
            try:
                if not DEBUG:
                    with open(file=grub_default, mode='a', encoding='utf-8') as f:
                        f.write(replacement + '\n')
                print(f"Appended '{replacement}' to {grub_default}")
            except IOError:
                print(f"Failed to append '{replacement}' to {grub_default}")
        else:
            print(f"No matches found for {pattern} in {grub_default}")

def main():
    if len(sys.argv) < 2:
        print("Usage: python3 themes_sudo.py <grub or plymouth> <optional: theme_name> <optional: grub_resolution>")
        sys.exit(1)

    theme_provider = sys.argv[1]
    theme_name = sys.argv[2] if len(sys.argv) >= 3 else None
    resolution = sys.argv[3] if len(sys.argv) >= 4 else None

    if theme_name is None:
        print(f"{theme_provider} theme name not provided")
    else:
        print(f"{theme_provider} theme name: {theme_name}")

    grub_default = '/etc/default/grub'
    set_theme_path = '/usr/sbin/plymouth-set-default-theme'
    modules_path = '/etc/initramfs-tools/modules'

    if theme_provider == 'plymouth':
        # Create empty modules file
        if not os.path.isfile(modules_path):
            try:
                with open(file=modules_path, mode='w', encoding='utf-8'):
                    pass
                print(f"Empty file created: {modules_path}")
            except IOError:
                modules_path = None

        clean_modules_file(modules_path=modules_path)

        if theme_name is not None:
            print(f"Set Plymouth theme: {set_theme_path} {theme_name}")
            if not DEBUG:
                try:
                    subprocess.run(
                        [set_theme_path, theme_name],
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        check=True
                    )
                except subprocess.CalledProcessError as e:
                    print(f"Failed to set Plymouth theme: {e}")
                    sys.exit(1)

            # Set Grub splash
            add_splash_to_grub(grub_default=grub_default, enable_splash=True)

            plymouth_update()

            print(f"Plymouth theme saved: {theme_name}")
        else:
            # Set Grub nosplash
            add_splash_to_grub(grub_default=grub_default, enable_splash=False)

            plymouth_update()

            print("Plymouth theme removed")
    elif theme_provider == "grub":
        if resolution is None:
            print(f"{theme_provider} resolution not provided")
        else:
            print(f"{theme_provider} resolution: {resolution}")

        # Themes can be in /boot/grub/themes or /usr/share/grub/themes
        theme_dir_boot = os.path.join('/boot/grub/themes', theme_name)
        theme_dir_user = os.path.join('/usr/share/grub/themes', theme_name)
        theme_dir = theme_dir_boot
        if os.path.isdir(theme_dir_user):
            theme_dir = theme_dir_user

        # Check if / is encrypted and /boot is not encrypted: copy theme to /boot/grub/themes
        if theme_dir != theme_dir_boot:
            if is_encrypted('/') and \
                not is_encrypted('/boot') and \
                os.path.isdir(theme_dir_user) and \
                not os.path.isdir(theme_dir_boot):
                # Copy the themes to /boot/grub/themes
                print("Copying themes to /boot/grub/themes")
                shutil.copytree(theme_dir_user, theme_dir_boot, dirs_exist_ok=True)
                theme_dir = theme_dir_boot

        theme_path = os.path.join(theme_dir, 'theme.txt')
        if not os.path.exists(theme_path):
            print(f"Grub theme not found: {theme_path}")
            sys.exit(1)

        print(f"Set Grub theme in {grub_default}: {theme_name} - {theme_path}")

        theme_bg = os.path.join(theme_dir, 'bg.png')

        if theme_path is not None:
            replace_pattern_in_grub(grub_default,
                                    r'^#*\s*GRUB_THEME\s*=.*$',
                                    f'GRUB_THEME={theme_path}')
            replace_pattern_in_grub(grub_default,
                                    r'^#*\s*GRUB_BACKGROUND\s*=.*$',
                                    f'GRUB_BACKGROUND={theme_bg}')
            print(f'Grub theme set: {theme_path}')
        else:
            replace_pattern_in_grub(grub_default,
                                    r'^GRUB_THEME\s*=',
                                    '#GRUB_THEME=',
                                    False)
            replace_pattern_in_grub(grub_default,
                                    r'^GRUB_BACKGROUND\s*=',
                                    '#GRUB_BACKGROUND=',
                                    False)
            print('Grub theme disabled')
            resolution = None

        if resolution is not None:
            replace_pattern_in_grub(grub_default,
                                r'^#*\s*GRUB_GFXMODE\s*=.*$',
                                f'GRUB_GFXMODE={resolution},auto')
            replace_pattern_in_grub(grub_default,
                                    r'^#*\s*GRUB_GFXPAYLOAD_LINUX\s*=.*$',
                                    'GRUB_GFXPAYLOAD_LINUX=keep')
            print(f'Grub resolution set: {resolution}')
        else:
            replace_pattern_in_grub(grub_default,
                                r'^GRUB_GFXMODE\s*=',
                                '#GRUB_GFXMODE=',
                                False)
            replace_pattern_in_grub(grub_default,
                                    r'^GRUB_GFXPAYLOAD_LINUX\s*=',
                                    '#GRUB_GFXPAYLOAD_LINUX=',
                                    False)
            print('Grub resolution disabled')

        # Run update-grub in a separate thread to prevent blocking
        run_command(['update-grub'])

if __name__ == "__main__":
    # Start logging
    print("themes_sudo.py started.")
    main()
