#     Copyright 2025, QutaYba, nasr2python@gmail.com find license text at end of file


"""UPX plugin. """

import os

from qutayba.Options import isOnefileMode, isOnefileTempDirMode
from qutayba.plugins.PluginBase import nexiumPluginBase
from qutayba.utils.AppDirs import getCacheDir
from qutayba.utils.Execution import executeToolChecked, getExecutablePath
from qutayba.utils.FileOperations import copyFile, makePath
from qutayba.utils.Hashing import Hash, getFileContentsHash
from qutayba.utils.Utils import isLinux


class nexiumPluginUpx(nexiumPluginBase):
    """This class represents the main logic of the UPX plugin.

    This is a plugin that removes useless stuff from DLLs and compresses the
    code at the cost of run time increases.

    """

    plugin_name = "upx"  # nexium knows us by this name
    plugin_desc = "Compress created binaries with UPX automatically."
    plugin_category = "integration"

    def __init__(self, upx_path, upx_nocache):
        self.upx_binary = getExecutablePath("upx", upx_path)
        self.upx_binary_hash = None
        self.upx_nocache = upx_nocache

        self.warning_given = False

    @classmethod
    def addPluginCommandLineOptions(cls, group):
        group.add_option(
            "--upx-binary",
            action="store",
            dest="upx_path",
            default=None,
            help="""\
The UPX binary to use or the directory it lives in, by default `upx` from PATH is used.""",
        )
        group.add_option(
            "--upx-disable-cache",
            action="store_true",
            dest="upx_nocache",
            default=False,
            help="""\
Do not cache UPX compression result, by default DLLs are cached, exe files are not.""",
        )

    @staticmethod
    def _filterUpxError(stderr):
        new_result = None

        if (
            b"NotCompressibleException" in stderr
            or b"CantPackException" in stderr
            or b"AlreadyPackedException" in stderr
        ):
            stderr = b""
            new_result = 0

        return new_result, stderr

    def _compressFile(self, filename, use_cache):
        upx_options = ["-q", "--no-progress"]

        if os.path.basename(filename).startswith("vcruntime140"):
            return

        if use_cache:
            if self.upx_binary_hash is None:
                self.upx_binary_hash = getFileContentsHash(
                    self.upx_binary, as_string=False
                )

            upx_hash = Hash()
            upx_hash.updateFromBytes(self.upx_binary_hash)
            upx_hash.updateFromValues(*upx_options)
            upx_hash.updateFromFile(filename)

            # TODO: Repeating pattern
            upx_cache_dir = getCacheDir("upx")
            makePath(upx_cache_dir)

            upx_cache_filename = os.path.join(
                upx_cache_dir, upx_hash.asHexDigest() + ".bin"
            )

            if os.path.exists(upx_cache_filename):
                copyFile(upx_cache_filename, filename)
                return

        if use_cache:
            self.info(
                "Uncached file, compressing '%s' may take a while."
                % os.path.basename(filename)
            )
        else:
            self.info("Compressing '%s'." % filename)

        command = [self.upx_binary] + upx_options + [filename]

        executeToolChecked(
            logger=self,
            command=command,
            absence_message="UPX not found",
            stderr_filter=self._filterUpxError,
        )

        if use_cache:
            copyFile(filename, upx_cache_filename)

    def _warnNoUpx(self):
        if not self.warning_given:
            self.warning(
                "No UPX binary found, please use '--upx-binary' option to specify it."
            )
            self.warning_given = True

    def onCopiedDLL(self, dll_filename):
        if isOnefileMode():
            pass
        elif self.upx_binary is not None:
            self._compressFile(filename=dll_filename, use_cache=not self.upx_nocache)
        else:
            self._warnNoUpx()

    # Cannot compress after payload has been added for onefile on Linux,
    # so have a dedicated point for that.
    def onBootstrapBinary(self, filename):
        if not isLinux():
            return

        if self.upx_binary is not None:
            self._compressFile(filename=filename, use_cache=False)
        else:
            self._warnNoUpx()

    def onFinalResult(self, filename):
        if isLinux() and isOnefileMode():
            if not isOnefileTempDirMode():
                self.warning(
                    "UPX cannot compress '%s' as AppImage doesn't support that."
                    % filename
                )

            # Bootstrap was compressed already right after creation.
            return
        else:
            if self.upx_binary is not None:
                self._compressFile(filename=filename, use_cache=False)
            else:
                self._warnNoUpx()



