#!/usr/bin/env python3
"""
Hybrid APT/PackageKit handler.

- PackageKit for: simple install/remove, cache refresh
- APT (via pkexec) for: updates, backports, third-party repos, system upgrades

Reference:
- https://lazka.github.io/pgi-docs/PackageKitGlib-1.0

Author: Arjen Balfoort
Date: 2026-04-18
"""

import os
import glob
import logging
import threading
import subprocess
import gettext
import queue
import re
from pathlib import Path
import psutil
import apt

from app_shared import (
    AppStreamPackageInfo,
    PackageAction,
    DebSoftwareApp,
    LogLevel,
)

import gi
gi.require_version("PackageKitGlib", "1.0")
from gi.repository import PackageKitGlib, GLib, Gio


_ = gettext.translation("deb-software", fallback=True).gettext


# =========================================
# PackageKit status/error maps
# =========================================

STATUS_MAP: dict[PackageKitGlib.StatusEnum, str] = {
    PackageKitGlib.StatusEnum.UNKNOWN: '',
    PackageKitGlib.StatusEnum.WAIT: _("Waiting"),
    PackageKitGlib.StatusEnum.UPDATE: _("Updating"),
    PackageKitGlib.StatusEnum.CLEANUP: _("Cleaning up"),
    PackageKitGlib.StatusEnum.DEP_RESOLVE: _("Resolving dependencies"),
    PackageKitGlib.StatusEnum.SIG_CHECK: _("Checking signature"),
    PackageKitGlib.StatusEnum.TEST_COMMIT: _("Testing commit"),
    PackageKitGlib.StatusEnum.COMMIT: _("Committing changes"),
    PackageKitGlib.StatusEnum.FINISHED: _("Finished"),
    PackageKitGlib.StatusEnum.CANCEL: _("Canceled"),
    PackageKitGlib.StatusEnum.SETUP: _("Preparing"),
    PackageKitGlib.StatusEnum.DOWNLOAD_REPOSITORY: _("Downloading repository information"),
    PackageKitGlib.StatusEnum.DOWNLOAD_PACKAGELIST: _("Downloading package list"),
    PackageKitGlib.StatusEnum.DOWNLOAD_FILELIST: _("Downloading file list"),
    PackageKitGlib.StatusEnum.DOWNLOAD_CHANGELOG: _("Downloading changelog"),
    PackageKitGlib.StatusEnum.DOWNLOAD_GROUP: _("Downloading group information"),
    PackageKitGlib.StatusEnum.DOWNLOAD_UPDATEINFO: _("Downloading update information"),
    PackageKitGlib.StatusEnum.REPACKAGING: _("Repackaging files"),
    PackageKitGlib.StatusEnum.LOADING_CACHE: _("Loading cache"),
    PackageKitGlib.StatusEnum.SCAN_APPLICATIONS: _("Scanning applications"),
    PackageKitGlib.StatusEnum.GENERATE_PACKAGE_LIST: _("Generating package list"),
    PackageKitGlib.StatusEnum.RUNNING: _("Running"),
    PackageKitGlib.StatusEnum.WAITING_FOR_LOCK: _("Waiting for package manager lock"),
    PackageKitGlib.StatusEnum.WAITING_FOR_AUTH: _("Waiting for authorization"),
    PackageKitGlib.StatusEnum.SCAN_PROCESS_LIST: _("Scanning process list"),
    PackageKitGlib.StatusEnum.CHECK_EXECUTABLE_FILES: _("Checking executable files"),
    PackageKitGlib.StatusEnum.CHECK_LIBRARIES: _("Checking libraries"),
    PackageKitGlib.StatusEnum.COPY_FILES: _("Copying files"),
    PackageKitGlib.StatusEnum.RUN_HOOK: _("Running package hook"),
    PackageKitGlib.StatusEnum.QUERY: _("Querying"),
    PackageKitGlib.StatusEnum.INFO: _("Retrieving information"),
    PackageKitGlib.StatusEnum.REMOVE: _("Removing"),
    PackageKitGlib.StatusEnum.REFRESH_CACHE: _("Updating cache"),
    PackageKitGlib.StatusEnum.DOWNLOAD: _("Downloading packages"),
    PackageKitGlib.StatusEnum.INSTALL: _("Installing"),
}

ERROR_MAP: dict[PackageKitGlib.ErrorEnum, str] = {
    PackageKitGlib.ErrorEnum.NO_NETWORK: _("No network connection."),
    PackageKitGlib.ErrorEnum.NOT_AUTHORIZED: _("Authorization failed."),
    PackageKitGlib.ErrorEnum.NO_CACHE: _("The package cache is not available."),
    PackageKitGlib.ErrorEnum.INTERNAL_ERROR: _("An internal error occurred."),
    PackageKitGlib.ErrorEnum.NOT_SUPPORTED: _("This operation is not supported."),
}


# =========================================
# APT high-level status/error mapping
# =========================================

def map_apt_status(line: str) -> str | None:
    """
    Map raw apt-get output lines to high-level, user-friendly messages.
    """
    l = line.lower()

    if "reading package lists" in l:
        return _("Reading package lists")
    if "building dependency tree" in l:
        return _("Resolving dependencies")
    if "reading state information" in l:
        return _("Checking installed packages")
    if "downloading" in l or "get:" in l or 'fetched' in l:
        return _("Downloading packages")
    if "preparing to unpack" in l:
        return _("Preparing to unpack packages")
    if "unpacking" in l:
        return _("Unpacking packages")
    if "setting up" in l:
        return _("Configuring packages")
    if "processing triggers" in l:
        return _("Processing triggers")
    if "done" == l.strip():
        return _("Finishing")

    return None


def map_apt_error(line: str) -> str | None:
    """
    Map apt-get error lines to high-level error messages.
    """
    l = line.lower()

    if "e:" in l and "unmet dependencies" in l:
        return _("Unmet dependencies prevent this operation.")
    if "e:" in l and "could not get lock" in l:
        return _("Another package manager is running.")
    if "e:" in l and "not found" in l:
        return _("One or more packages were not found.")
    if "e:" in l and "permission denied" in l:
        return _("Permission denied while running APT.")
    if "e:" in l:
        return _("APT reported an error.")

    return None


# =========================================
# Base handler
# =========================================

class BaseHandler(threading.Thread):
    """
    Base class with shared queue/logging helpers.
    """

    def __init__(self,
                 app: DebSoftwareApp,
                 action: PackageAction,
                 packages: list[AppStreamPackageInfo] | None = None,
                 result_queue: queue.Queue | None = None,
                 on_done=None) -> None:
        super().__init__(daemon=True)
        self.app: DebSoftwareApp = app
        self.script_dir: str = app.script_dir
        self.logging: logging = app.logging
        self.result_queue: queue.Queue = result_queue
        self.packages: list[AppStreamPackageInfo] = packages or []
        self.action: PackageAction = action
        self.on_done: callable = on_done
        self._process: subprocess.Popen[str] | None = None
        special_scripts = glob.glob(os.path.join(self.script_dir, 'scripts/*.sh'))
        self.specials: list[str] = []
        for script in special_scripts:
            self.specials.append(Path(script).stem)

    def queue_message(self,
                      level: LogLevel,
                      message: str,
                      percentage: int) -> None:
        """
        Add a message to the queue with its level (info/error).
        """
        if self.result_queue is not None:
            info_item = {
                'level': level,
                'message': message,
                'progress': percentage
            }
            self.result_queue.put(info_item)

        if self.logging is not None and message:
            message = message.replace('\n', ' ')
            if level in (LogLevel.INFO, LogLevel.SUCCESS, LogLevel.PROGRESS):
                self.logging.info(message)
            elif level in (LogLevel.WARNING, LogLevel.CANCELED):
                self.logging.warning(message)
            elif level == LogLevel.ERROR:
                self.logging.error(message)
            elif level == LogLevel.DEBUG:
                self.logging.debug(message)

    @staticmethod
    def _extract_percentage(line: str) -> int:
        match = re.search(r"(\d+)%", line)
        return int(match.group(1)) if match else 0

    @staticmethod
    def _extract_num_pkgs(line: str) -> tuple[bool, int]:
        upgraded = re.search(r"(\d+)\s+upgraded", line)
        newly = re.search(r"(\d+)\s+newly\s+installed", line)

        if upgraded or newly:
            up = int(upgraded.group(1)) if upgraded else 0
            new = int(newly.group(1)) if newly else 0
            return True, max(1, up + new)

        return False, 1

    @staticmethod
    def _apt_is_locked() -> bool:
        """Check if dpkg is locked."""
        filepath = "/var/lib/dpkg/lock"
        for proc in psutil.process_iter(['pid', 'name', 'open_files']):
            try:
                for file in proc.info['open_files'] or []:
                    if filepath in file.path:
                        return True
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                continue
        return False

    # ---- APT subprocess runner -----------------------------------------

    def _run_apt_command(self, command: list[str]) -> None:
        if self._apt_is_locked():
            self.queue_message(
                LogLevel.ERROR,
                _("APT is currently busy.\nPlease wait for the ongoing operation to finish."),
                0
            )
            return
        self.logging.debug(f"Running APT command: {' '.join(command)}")

        # Check if this is a flatpak install
        flatpak_install = any("flatpak" in l for l in command)

        with subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
        ) as process:
            self._process = process

            num_pkgs = 1
            cur_pkg = 1
            num_pkgs_done = False
            last_global_pct = 0
            seen_get_lines = 0

            try:
                for line in process.stdout:
                    line = line.rstrip()
                    if not line:
                        continue

                    if not num_pkgs_done:
                        num_pkgs_done, num_pkgs = self._extract_num_pkgs(line)

                    # Download progress
                    if "Get:" in line:
                        seen_get_lines += 1
                        if num_pkgs > 0:
                            pct = int((seen_get_lines / num_pkgs) * 100)
                            status_msg = map_apt_status(line)
                            if status_msg:
                                self.queue_message(LogLevel.PROGRESS, status_msg, pct)
                            else:
                                self.queue_message(LogLevel.DEBUG, line, pct)
                        continue

                    # Upacking progress
                    if "Preparing to unpack" in line:
                        cur_pkg += 1

                    extracted_pct = self._extract_percentage(line)
                    global_pct = int(
                        ((cur_pkg - 1) + extracted_pct / 100.0) / num_pkgs * 100
                    )
                    global_pct = min(max(global_pct, last_global_pct), 100)
                    last_global_pct = global_pct

                    status_msg = map_apt_status(line)
                    if status_msg:
                        self.queue_message(LogLevel.PROGRESS, status_msg, global_pct)
                    else:
                        self.queue_message(LogLevel.DEBUG, line, global_pct)

                try:
                    returncode = process.wait(timeout=3600)
                except subprocess.TimeoutExpired:
                    self.logging.error("APT command timed out")
                    self.queue_message(
                        LogLevel.ERROR,
                        _("APT operation timed out."),
                        last_global_pct,
                    )
                    return

                # Only show error messages if the return code is not 127 (user cancel)
                if returncode != 127:
                    for line in process.stderr:
                        line = line.rstrip()
                        if not line:
                            continue

                        err_msg = map_apt_error(line)
                        if err_msg:
                            self.queue_message(LogLevel.ERROR, err_msg, 0)
                        else:
                            self.queue_message(
                                LogLevel.DEBUG, f"apt stderr: {line}", 0
                            )
            finally:
                self._process = None

        if returncode != 0:
            # user cancel
            if returncode not in (126, 127):
                self.queue_message(
                    LogLevel.ERROR,
                    _("APT failed with exit code %s") % returncode,
                    last_global_pct,
                )
            else:
                self.queue_message(
                    LogLevel.CANCELED,
                    _("Canceled by user"),
                    last_global_pct,
                )
        else:
            # Update AppStream when flatpak was installed
            if flatpak_install:
                command = ['flatpak', 'update', '--appstream']
                self.logging.debug(f"Running command: {' '.join(command)}")
                self.queue_message(LogLevel.INFO,
                                _("Updating AppStream"), last_global_pct)
                subprocess.run(command, check=True)

            self.queue_message(
                LogLevel.SUCCESS,
                _("APT operation completed."),
                100,
            )

# =========================================
# PackageKit handler
# =========================================

class PackageKitHandler(BaseHandler):
    """
    Uses PackageKit for:
    - install
    - remove
    - refresh cache
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.client: PackageKitGlib.Client = PackageKitGlib.Client()
        self.package_ids: list[str] = []
        if self.action != PackageAction.REFRESH_CACHE:
            self.package_ids: list[str] = self.search_package_ids()

    def run(self) -> None:
        try:
            if self.action == PackageAction.INSTALL:
                # Special cases need scripts
                self._check_for_specials()

                # Normal packages
                if self.package_ids:
                    self._install_packages()

            elif self.action == PackageAction.REMOVE:
                self._remove_packages()
            elif self.action == PackageAction.REFRESH_CACHE:
                self._refresh_cache()
            else:
                self.queue_message(LogLevel.ERROR,
                    _("Unsupported PackageKit action") + f": {self.action.value}", 0)
        finally:
            if self.on_done:
                self.on_done()

    def search_package_ids(self) -> list[str]:
        """ Searches for package ids. """
        ids: list[str] = []
        for pkg in self.packages:
            try:
                results = self.client.search_names(
                    PackageKitGlib.FilterEnum.NONE,
                    [pkg.package_name],
                    None,
                    self._progress_callback,
                    None,
                )
            except GLib.Error as e:
                self.queue_message(LogLevel.ERROR, str(e), 0)
                continue

            for p in results.get_package_array():
                if p.get_name() == pkg.package_name:
                    # Flatpak is a special case
                    if pkg.package_name not in self.specials:
                        ids.append(p.get_id())
                    break

        return ids

    def _check_for_specials(self) -> None:
        """
        Some packages need some special handling.
        Bash scripts are located in /usr/lib/deb-software/scripts

        Args:
            scripts_paths (list[str]): List of script paths to run
        """
        special_scripts = [
            os.path.join(self.script_dir, 'scripts', f'{special}.sh')
            for special in self.specials
            if special in {pkg.package_name for pkg in self.packages}
        ]

        for script_path in special_scripts:
            command = [
                'pkexec',
                '--disable-internal-agent',
                '/usr/lib/deb-software/scripts/_run_script.sh',
                script_path
            ]

            self._run_apt_command(command)

    def _new_cancellable(self):
        return Gio.Cancellable()

    def _progress_callback(
        self,
        progress: PackageKitGlib.Progress,
        progress_type: PackageKitGlib.ProgressType,
        _user_data=None,
    ) -> None:
        if progress_type == PackageKitGlib.ProgressType.PERCENTAGE:
            pct = max(0, progress.get_percentage())
            status = progress.get_status()
            message = STATUS_MAP.get(status, "")
            self.queue_message(LogLevel.PROGRESS, message, pct)

    def _install_packages(self) -> None:
        try:
            msg = self._status_message() + " " + \
                ', '.join([pkg.package_name for pkg in self.packages])
            self.queue_message(LogLevel.PROGRESS, msg, 0)
            results = self.client.install_packages(
                PackageKitGlib.TransactionFlagEnum.NONE,
                self.package_ids,
                self._new_cancellable(),
                self._progress_callback,
                None,
            )
            self._store_results(results)
        except GLib.Error as e:
            self._store_error(e)

    def _remove_packages(self) -> None:
        try:
            msg = self._status_message() + " " + \
                ', '.join([pkg.package_name for pkg in self.packages])
            self.queue_message(LogLevel.PROGRESS, msg, 0)
            results = self.client.remove_packages(
                PackageKitGlib.TransactionFlagEnum.NONE,
                self.package_ids,
                True, # remove dependencies
                True, # autoremove
                self._new_cancellable(),
                self._progress_callback,
                None,
            )
            self._store_results(results)
        except GLib.Error as e:
            self._store_error(e)

    def _refresh_cache(self) -> None:
        try:
            self.queue_message(LogLevel.PROGRESS, self._status_message(), 0)
            results = self.client.refresh_cache(
                True, # force: aggresively drop caches
                self._new_cancellable(),
                self._progress_callback,
                None,
            )
            self._store_results(results)
        except GLib.Error as e:
            self._store_error(e)

    def _store_results(self, results) -> None:
        exit_status = results.get_exit_code()
        error = results.get_error_code()

        if exit_status == PackageKitGlib.ExitEnum.SUCCESS:
            pkgs = ', '.join([pkg.package_name for pkg in self.packages])
            msg = _("Finished") + f" {self._status_message().lower()} {pkgs}"
            self.queue_message(LogLevel.SUCCESS,
                               msg,
                               100)
        elif error is not None:
            error_code = error.get_code()
            error_details = error.get_details()
            msg = ERROR_MAP.get(error_code, error_details)
            self.queue_message(LogLevel.ERROR, msg, 0)

    def _store_error(self, error: GLib.Error) -> None:
        # Skip user not providing sudo password
        if error.code == 303:
            self.queue_message(LogLevel.CANCELED, _("Canceled by user"), 0)
            return
        # Try to map PackageKit error to friendly text
        try:
            msg = ERROR_MAP.get(error.code, str(error))
        except KeyError:
            msg = str(error)
        self.queue_message(LogLevel.ERROR, msg, 0)

    def _status_message(self) -> str:
        if self.action == PackageAction.INSTALL:
            return _("Installing")
        if self.action == PackageAction.REMOVE:
            return _("Removing")
        if self.action == PackageAction.UPDATE:
            return _("Updating")
        if self.action == PackageAction.REFRESH_CACHE:
            return _("Updating cache")
        return _("Working")


# =========================================
# APT subprocess handler (via pkexec)
# =========================================

class AptSubprocessHandler(BaseHandler):
    """
    Uses apt-get via pkexec for:
    - updating specific packages
    - full system upgrades
    - backports / third-party repos
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.apt_cache: apt.Cache = apt.Cache()

    def run(self) -> None:
        try:
            if self.action == PackageAction.UPDATE:
                self._apt_update_packages()
            elif self.action == PackageAction.SYSTEM_UPGRADE:
                self._apt_dist_upgrade()
            else:
                self.queue_message(LogLevel.ERROR, _("Unsupported APT action."), 0)
        finally:
            if self.on_done:
                self.on_done()

    def cancel(self) -> None:
        """
        Cancels the APT operation.
        Hook cancel() to a “Cancel” button.

        Currently not used: PackageKit and Flatpak not implemented, yet
        """
        proc = self._process
        if proc and proc.poll() is None:
            self.logging.info("Cancelling APT operation...")
            try:
                proc.terminate()
            except Exception:
                pass
            try:
                proc.kill()
            except Exception:
                pass

    # ---- APT operations ------------------------------------------------

    def _apt_update_packages(self) -> None:
        pkg_names = [pkg.package_name for pkg in self.packages]
        if not pkg_names:
            self.queue_message(LogLevel.ERROR, _("No packages to update."), 0)
            return

        # pkexec will prompt for root password via polkit
        command = ["pkexec",
                   "--disable-internal-agent",
                   f"{self.script_dir}/scripts/_apt_install.sh"] + pkg_names
        self.queue_message(LogLevel.PROGRESS, _("Updating selected packages"), 0)
        self._run_apt_command(command)

    def _apt_dist_upgrade(self) -> None:
        # pkexec will prompt for root password via polkit
        command = ["pkexec",
                   "--disable-internal-agent",
                   f"{self.script_dir}/scripts/_dist_upgrade.sh"]
        self.queue_message(LogLevel.PROGRESS, _("Upgrading system"), 0)
        self._run_apt_command(command)

# =========================================
# Hybrid handler - start here
# =========================================

class AptPackageHandler(BaseHandler):
    """
    Hybrid handler that chooses between PackageKit and APT.

    - PackageKit:
        * INSTALL
        * REMOVE
        * REFRESH_CACHE

    - APT (via pkexec):
        * UPDATE
        * SYSTEM_UPGRADE
    """

    def run(self) -> None:
        try:
            if self._should_use_packagekit():
                self._run_packagekit()
            else:
                self._run_apt()
        finally:
            if self.on_done:
                self.on_done()

    def _should_use_packagekit(self) -> bool:
        """
        Decide which backend to use.

        - Simple actions > PackageKit
        - Updates/system upgrades > APT
        """
        if self.action in (
            PackageAction.INSTALL,
            PackageAction.REMOVE,
            PackageAction.REFRESH_CACHE,
        ):
            return True
        # Use APT for updates and system upgrades
        return False

    def _run_packagekit(self) -> None:
        handler = PackageKitHandler(
            app=self.app,
            action=self.action,
            packages=self.packages,
            result_queue=self.result_queue,
            on_done=None,  # outer handler calls on_done
        )
        handler.run()

    def _run_apt(self) -> None:
        handler = AptSubprocessHandler(
            app=self.app,
            action=self.action,
            packages=self.packages,
            result_queue=self.result_queue,
            on_done=None,  # outer handler calls on_done
        )
        handler.run()
