#     Copyright 2021, Kay Hayen, mailto:kay.hayen@gmail.com
#
#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.
#
"""UPX plugin. """

import os

from nuitka.Options import isOnefileMode, isOnefileTempDirMode
from nuitka.plugins.PluginBase import NuitkaPluginBase
from nuitka.utils.AppDirs import getCacheDir
from nuitka.utils.Execution import check_call, getExecutablePath, getNullOutput
from nuitka.utils.FileOperations import copyFile, makePath
from nuitka.utils.Hashing import Hash, getFileContentsHash
from nuitka.utils.Utils import isLinux


class NuitkaPluginUpx(NuitkaPluginBase):
    """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"  # Nuitka knows us by this name
    plugin_desc = "Compress created binaries with UPX automatically"

    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.""",
        )

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

        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 = os.path.join(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)

        check_call(
            [self.upx_binary] + upx_options + [filename],
            stdout=getNullOutput(),
            shell=False,
        )

        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()
