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

Author: Arjen Balfoort 
Date: 2026-01-16
Dependencies: plymouth, (grub-common, apt, dpkg)
"""

import re
import os
import logging
import gettext
import subprocess
import shutil
from enum import StrEnum
from math import gcd
from typing import Optional

_ = gettext.translation('solydxk-system', fallback=True).gettext

class Themes(StrEnum):
    """
    Enum for themes
    """
    GRUB = 'grub'
    PLYMOUTH = 'plymouth'


class Plymouth():
    """
    Handles general plymouth functions
    """
    def __init__(self, logging_handler):
        self._set_theme_path = '/usr/sbin/plymouth-set-default-theme'
        self.is_plymouth_installed: bool = True
        self.logging_handler: logging = logging_handler

        if not os.path.isfile(self._set_theme_path):
            self.logging_handler.warning("%s not found - install plymouth and try again",
                                         self._set_theme_path)
            self.is_plymouth_installed = False

        # Get Plymouth themes (after LOGGING is set)
        self.installed_themes: Optional[list[str]] = self._get_installed_themes()

    @staticmethod
    def _filter_packages(packages: list[str]) -> list[str]:
        """
        Filters packages based on not installed status.

        Args:
            packages (list): A list of packages to filter.

        Returns:
            list: A list of packages not yet installed.
        """
        filtered_packages = []

        for package in packages:
            try:
                pck_result = subprocess.run(
                    ["dpkg", "-l", package],
                    capture_output=True,
                    text=True,
                    check=False
                ).stdout.strip()
                if "ii" not in pck_result:
                    filtered_packages.append(package)
            except subprocess.CalledProcessError:
                continue

        return filtered_packages

    def _get_installed_themes(self) -> Optional[list[str]]:
        """
        Returns a list of installed Plymouth themes.

        If the plymouth-set-default-theme command is available, it runs the command
        with the '--list' option and returns the output as a list of theme names.
        Otherwise, it returns an empty list.
        """
        if self.is_plymouth_installed:
            try:
                themes = subprocess.run(
                    [self._set_theme_path, '--list'],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    check=True,
                    text=True
                ).stdout.splitlines()
            except subprocess.CalledProcessError:
                return None

            # Remove duplicates (case-insensitive) and sort
            seen = set()
            unique_sorted = sorted(
                [s for s in themes if not (s.lower() in seen or seen.add(s.lower()))],
                key=lambda x: x.lower()
            )

            # Put solydxk themes first
            solydxk_themes = []
            other_themes = []
            for theme in unique_sorted:
                if "solyd" in theme.lower():
                    solydxk_themes.append(theme)
                else:
                    other_themes.append(theme)

            self.logging_handler.debug(
                "Installed Plymouth themes: %s", ', '.join(solydxk_themes + other_themes)
            )
            return solydxk_themes + other_themes

        return None

    def is_plymouth_booted(self) -> bool:
        """
        Returns True if Plymouth is booted, False otherwise.
        """
        try:
            cmdline = subprocess.run(
                ['cat', '/proc/cmdline'],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
                text=True
            ).stdout.strip()
        except subprocess.CalledProcessError:
            self.logging_handler.error("Unable to get /proc/cmdline contents")
            return False

        if ' splash' in cmdline:
            return True
        return False

    def get_current_theme(self) -> Optional[str]:
        """
        Returns the currently used Plymouth theme.
        """
        if self.is_plymouth_installed and \
           self.is_plymouth_booted():
            try:
                return subprocess.run(
                    [self._set_theme_path],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    check=True,
                    text=True
                ).stdout.strip()
            except subprocess.CalledProcessError:
                self.logging_handler.error("Unable to get current Plymouth theme")
                return None
        return None

    def get_available_theme_packages(self) -> Optional[list[str]]:
        """
        Returns a list of available Plymouth theme packages in the repositories.        
        """
        try:
            packages = subprocess.run(
                ['apt-cache', 'pkgnames', 'plymouth-theme'],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                check=True,
                text=True
            ).stdout.splitlines()
        except subprocess.CalledProcessError:
            return None

        # Remove packages that are already installed
        filtered_packages = self._filter_packages(packages)

        self.logging_handler.debug("Available Plymouth theme packages: %s", filtered_packages)
        return filtered_packages

    def preview_plymouth(self, in_tty8: bool = False) -> None:
        """
        Preview Plymouth themes with updates.

        Args:
            in_tty8 (bool, optional): If True, preview in tty8. Defaults to False.
        """
        debug_tty = '--debug --debug-file=/dev/tty8' if in_tty8 else ''
        command = f"""
su -c '
plymouthd {debug_tty};
plymouth --show-splash;
for ((I=0; I<10; I++)); do
    plymouth --update=test$I;
    sleep 1;
done;
plymouth quit
'
"""
        try:
            subprocess.run(command, shell=True, check=True, executable="/bin/bash")
            self.logging_handler.info("Plymouth theme preview with updates completed.")
            self.logging_handler.info("Switch to tty8 (Ctrl+Alt+F8) to see the splash screen.")
        except subprocess.CalledProcessError as e:
            self.logging_handler.error("An error occurred: %s", e)

    def get_removable_package_name(self, theme_name: str) -> Optional[str]:
        """
        Get the package name that can be uninstalled of a given Plymouth theme.

        Args:
            theme_name (str): The name of the Plymouth theme.

        Returns:
            str: The package name that can be uninstalled, or None if not found.
        """
        package = None
        try:
            package_names = subprocess.run(
                    ['dpkg', '-S', theme_name + '.plymouth'],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    check=True,
                    text=True
                ).stdout.splitlines()
        except subprocess.CalledProcessError:
            self.logging_handler.error("Plymouth theme package %s not found", theme_name)
            return None

        for line in package_names:
            if 'plymouth-theme' in line:
                match = re.search('(^.*):', line)
                if match:
                    package = match.group(1)
                    break

        # Check if package is installed
        if not self._filter_packages([package]):
            return None

        self.logging_handler.debug("Removable Plymouth theme package found %s", package)
        return package

class Grub():
    """
    Manage Grub configuration.
    """
    def __init__(self, logging_handler):
        self.logging_handler = logging_handler
        self._grub_default = '/etc/default/grub'
        # Get Grub themes (after LOGGING is set)
        self.installed_themes = self._get_installed_themes()

    def _grub_default_value(self, pattern: str, group_nr: int=0) -> Optional[str]:
        """
        Get a value from the Grub configuration file.

        Args:
            pattern (str): The regex pattern to search for.
            group_nr (int, optional): The group number to extract from the match. Default is 0.

        Returns:
            Optional[str]: The value if found, None otherwise.
        """
        if not os.path.isfile('/etc/default/grub'):
            if os.path.exists('/usr/share/grub/default/grub'):
                shutil.copy('/usr/share/grub/default/grub', '/etc/default/grub')
            else:
                self.logging_handler.error("Grub configuration file not found")
                return None

        lines = []
        with open(file=self._grub_default, mode='r', encoding='utf-8') as grub_fle:
            lines = grub_fle.read().splitlines()
        for line in lines:
            # Search text for resolution
            match = re.search(pattern, line)
            if match:
                return match.group(group_nr).strip()
        return None

    def _get_installed_themes(self) -> list[tuple[str, str]]:
        """
        Get a list of installed Grub themes.

        Returns:
            list: A list of installed Grub themes and their paths.
        """
        themes = []
        for theme_dir in ('/boot/grub/themes', '/usr/share/grub/themes'):
            if not os.path.exists(theme_dir):
                continue

            for theme in set(next(os.walk(theme_dir))[1]):
                theme_path = os.path.join(theme_dir, theme + "/theme.txt")
                if os.path.exists(theme_path):
                    # Get theme name from path
                    theme_name = os.path.basename(os.path.dirname(theme_path))
                    themes.append(theme_name)

        # Remove duplicates and sort
        themes = sorted(list(set(themes)))

        # Put solydxk themes first
        solydxk_themes = []
        other_themes = []
        for theme in themes:
            if "solyd" in theme.lower():
                solydxk_themes.append(theme)
            else:
                other_themes.append(theme)

        # Sort on first element in tuple
        self.logging_handler.debug("Installed Grub themes: %s",
                                   ', '.join(solydxk_themes + other_themes))
        return solydxk_themes + other_themes

    def get_current_theme(self) -> Optional[tuple[str, str]]:
        """
        Get the current Grub theme and path.
        """
        theme_path = self._grub_default_value(r'^GRUB_THEME\s*=(.*)', 1)
        if os.path.exists(str(theme_path)):
            theme_name = os.path.basename(os.path.dirname(theme_path))
            self.logging_handler.debug("Current Grub theme: %s", theme_name)
            return theme_name
        return None

    def get_theme_path_from_name(self, theme_name: str) -> Optional[str]:
        """
        Get the path of the theme from its name.
        """
        for theme in self._get_installed_themes():
            if theme_name in theme:
                self.logging_handler.debug("Found Grub theme %s", theme)
                return theme
        return None

    @staticmethod
    def _simplify_ratio(width: int, height: int) -> str:
        """Simplify aspect ratio to human-readable format (e.g., 16:9)."""
        divisor = gcd(width, height)
        return f"{width // divisor}:{height // divisor}"

    def get_all_resolutions(self) -> list[tuple[str, str]]:
        """Get all resolutions from xrandr, not just the connected ones."""
        try:
            result = subprocess.run(['xrandr'],
                                    capture_output=True,
                                    text=True,
                                    check=True)
            output = result.stdout
            resolutions = []
            for line in output.splitlines():
                # Find all resolutions in the line
                matches = re.findall(r'(\d+x\d+)\s+\d+\.\d+\*?', line)
                for resolution in matches:
                    width, height = map(int, resolution.split('x'))
                    if width in [3840, 2560, 1920, 1600, 1366, 1280, 1024, 800, 640]:
                        aspect_ratio = self._simplify_ratio(width, height)
                        if aspect_ratio in ['4:3', '16:9']:
                            resolutions.append((resolution, aspect_ratio))
            # Sort by aspect ratio (16:9 first), then by width (high to low)
            resolutions.sort(key=lambda x: (0 if x[1] == '16:9' else 1,
                             -int(x[0].split('x')[0]),
                             -int(x[0].split('x')[1])))

            self.logging_handler.debug("Available Grub resolutions: %s", resolutions)
            return resolutions
        except FileNotFoundError:
            return []

    # Get current Grub resolution
    def get_current_resolution(self) -> Optional[tuple[str, str]]:
        """
        Get the current Grub resolution.
        """
        resolution = self._grub_default_value(r'^GRUB_GFXMODE\s*=[\s"]*([0-9]+x[0-9]+)', 1)
        if resolution is None or not 'x' in resolution:
            return (None, None)
        width, height = map(int, resolution.split('x'))
        aspect_ratio = self._simplify_ratio(width, height)

        self.logging_handler.debug("Current Grub resolution: %s (%s)", resolution, aspect_ratio)
        return (resolution, aspect_ratio)
