#     Copyright 2026, Kay Hayen, mailto:kay.hayen@gmail.com find license text at end of file


"""Internal tool, compression standalone distribution files and attach to onefile bootstrap binary."""

import os
import shutil
import struct
import sys
from contextlib import contextmanager

from nuitka.__past__ import to_byte
from nuitka.Progress import (
    closeProgressBar,
    enableProgressBar,
    reportProgressBar,
    setupProgressBar,
)
from nuitka.Tracing import onefile_logger
from nuitka.utils.AppDirs import getCacheDir
from nuitka.utils.FileOperations import areSamePaths, getFileList, getFileSize
from nuitka.utils.Hashing import Hash, HashCRC32
from nuitka.utils.Json import loadJsonFromFilename
from nuitka.utils.Utils import (
    decoratorRetries,
    getArchitecture,
    isWin32OrPosixWindows,
    isWin32Windows,
)
from nuitka.Version import version_string


def getCompressorLevel(low_memory):
    return 3 if low_memory else 22


def getCompressorFunction(expect_compression, low_memory, job_limit):
    # spell-checker: ignore closefd

    # On 32-bit Python, zstandard stream_writer at level 22 with threads >= 1
    # fails with "Allocation error : not enough memory" even for tiny data,
    # so we must use single-threaded compression on 32-bit.
    effective_threads = 0 if low_memory or getArchitecture() == "x86" else job_limit

    if expect_compression:
        try:
            from compression import zstd

            @contextmanager
            def useCompressedFile(output_file):
                options = {
                    zstd.CompressionParameter.nb_workers: effective_threads,
                    zstd.CompressionParameter.compression_level: getCompressorLevel(
                        low_memory
                    ),
                }

                with zstd.open(
                    output_file,
                    mode="wb",
                    options=options,
                ) as compressed_file:
                    yield compressed_file

        except ImportError:
            from zstandard import ZstdCompressor  # pylint: disable=I0021,import-error

            compressor_context = ZstdCompressor(
                level=getCompressorLevel(low_memory), threads=effective_threads
            )

            @contextmanager
            def useCompressedFile(output_file):
                with compressor_context.stream_writer(
                    output_file, closefd=False
                ) as compressed_file:
                    yield compressed_file

        onefile_logger.info("Using compression for onefile payload.")

        return b"Y", useCompressedFile
    else:

        @contextmanager
        def useSameFile(output_file):
            yield output_file

        return b"X", useSameFile


def _attachOnefilePayloadFile(
    output_file,
    is_archive,
    is_compressing,
    use_compression_cache,
    low_memory,
    file_compressor,
    filename_full,
    count,
    dist_dir,
    filename_encoding,
    file_checksums,
    win_path_sep,
):
    # Somewhat detail rich, at least unless we make more things mandatory, and
    # we also need to pass all modes, since this can be run in a separate process
    # that has no access to options.
    # pylint: disable=too-many-arguments,too-many-locals,too-many-statements

    payload_item_size = 0

    filename_relative = os.path.relpath(filename_full, dist_dir)

    reportProgressBar(
        item=filename_relative,
        update=False,
    )

    # Might be changing from POSIX to Win32 Python on Windows.
    if win_path_sep:
        filename_relative = filename_relative.replace("/", "\\")
    else:
        filename_relative = filename_relative.replace("\\", "/")

    filename_encoded = (filename_relative + "\0").encode(filename_encoding)

    output_file.write(filename_encoded)
    payload_item_size += len(filename_encoded)

    file_flags = 0
    if not isWin32OrPosixWindows() and os.path.islink(filename_full):
        link_target = os.readlink(filename_full)

        file_flags |= 2
        file_header = to_byte(file_flags)

        output_file.write(file_header)
        payload_item_size += len(file_header)

        link_target_encoded = (link_target + "\0").encode(filename_encoding)

        output_file.write(link_target_encoded)
        payload_item_size += len(link_target_encoded)
    else:
        # This flag is only relevant for non-links.
        if not isWin32OrPosixWindows() and os.access(filename_full, os.X_OK):
            file_flags |= 1

        with open(filename_full, "rb") as input_file:
            input_file.seek(0, 2)
            input_size = input_file.tell()
            input_file.seek(0, 0)

            file_header = b""

            if not isWin32OrPosixWindows():
                file_header += to_byte(file_flags)

            file_header += struct.pack("Q", input_size)

            if file_checksums:
                hash_crc32 = HashCRC32()
                hash_crc32.updateFromFileHandle(input_file)
                input_file.seek(0, 0)

                # CRC32 value 0 is avoided, used as error indicator in C code.
                file_header += struct.pack("I", hash_crc32.asDigest() or 1)

            if is_archive and is_compressing:
                compression_cache_filename = _getCacheFilename(
                    binary_filename=filename_full, low_memory=low_memory
                )

                if not os.path.exists(compression_cache_filename):
                    with open(compression_cache_filename, "wb") as archive_entry_file:
                        with file_compressor(
                            archive_entry_file
                        ) as compressed_file_tmp2:
                            shutil.copyfileobj(input_file, compressed_file_tmp2)

                compressed_size = getFileSize(compression_cache_filename)

                file_header += struct.pack("I", compressed_size)

            output_file.write(file_header)
            payload_item_size += len(file_header)

            if is_archive and is_compressing:
                with open(compression_cache_filename, "rb") as archive_entry_file:
                    pos1 = output_file.tell()
                    shutil.copyfileobj(archive_entry_file, output_file)
                    pos2 = output_file.tell()
                    assert pos2 - pos1 == compressed_size

                if count == 0 or not use_compression_cache:
                    os.unlink(compression_cache_filename)

                payload_item_size += compressed_size

            else:
                shutil.copyfileobj(input_file, output_file)
                payload_item_size += input_size

    reportProgressBar(
        item=filename_relative,
        update=True,
    )

    return payload_item_size


def _getCacheFilename(binary_filename, low_memory):
    hash_value = Hash()

    hash_value.updateFromFile(filename=binary_filename)

    # Have different values for different Python major versions.
    hash_value.updateFromValues(sys.version, sys.executable)

    # Take Nuitka version into account as well, ought to catch code changes.
    hash_value.updateFromValues(version_string)

    # Take zstandard version and compression level into account.
    try:
        from compression.zstd import zstd_version_info

        version = str(zstd_version_info)
    except ImportError:
        from zstandard import (
            __version__ as version,
        )  # pylint: disable=I0021,import-error

    hash_value.updateFromValues(version, getCompressorLevel(low_memory))

    cache_dir = getCacheDir("onefile-compression", create=True)

    return os.path.join(cache_dir, hash_value.asHexDigest())


def _getInputFileList(dist_dir, start_binary, include_start_binary, expected_files):
    file_list = getFileList(dist_dir, normalize=False)
    file_list_size = len(file_list)
    file_list = [
        filename for filename in file_list if not areSamePaths(filename, start_binary)
    ]

    if expected_files is not None:
        expected_files_set = set(
            os.path.normcase(os.path.normpath(f)) for f in expected_files
        )

        start_binary_relative_norm = os.path.normcase(
            os.path.normpath(os.path.relpath(start_binary, dist_dir))
        )
        if start_binary_relative_norm in expected_files_set:
            expected_files_set.remove(start_binary_relative_norm)

        for filename in file_list:
            filename_relative = os.path.relpath(filename, dist_dir)
            filename_relative_norm = os.path.normcase(
                os.path.normpath(filename_relative)
            )

            if filename_relative_norm not in expected_files_set:
                onefile_logger.warning(
                    """\
Distribution folder contains unexpected file '%s', this file will be ignored in future release."""
                    % filename_relative
                )
            else:
                expected_files_set.remove(filename_relative_norm)

        for missing_file in expected_files_set:
            onefile_logger.warning(
                """\
Distribution folder is missing expected file '%s'. This might cause issues."""
                % missing_file
            )

    if include_start_binary:
        file_list.insert(0, start_binary)

    # If this fails, above code failed to find the start binary.
    assert file_list_size == len(file_list) + (0 if include_start_binary else 1)

    return tuple(file_list)


def attachOnefilePayload(  # pylint: disable=too-many-arguments
    dist_dir,
    onefile_output_filename,
    start_binary,
    include_start_binary,
    expect_compression,
    as_archive,
    use_compression_cache,
    file_checksums,
    win_path_sep,
    low_memory,
    job_limit,
    expected_files,
):
    compression_indicator, compressor = getCompressorFunction(
        expect_compression=expect_compression,
        low_memory=low_memory,
        job_limit=job_limit,
    )

    @decoratorRetries(
        logger=onefile_logger,
        purpose="write payload to '%s'" % onefile_output_filename,
        consequence="the result is unusable",
    )
    def _attachOnefilePayload():
        with open(onefile_output_filename, "ab") as output_file:
            # Seeking to end of file seems necessary on Python2 at least, maybe it's
            # just that tell reports wrong value initially.
            output_file.seek(0, 2)
            start_pos = output_file.tell()
            output_file.write(b"KA" + compression_indicator)

            # Move the binary to start immediately to the start position
            file_list = _getInputFileList(
                dist_dir=dist_dir,
                start_binary=start_binary,
                include_start_binary=include_start_binary,
                expected_files=expected_files,
            )

            if isWin32Windows():
                filename_encoding = "utf-16le"
            else:
                filename_encoding = "utf8"

            payload_size = 0

            setupProgressBar(
                stage="Onefile Payload",
                unit="module",
                total=len(file_list),
            )

            # Abstract the differences here for the time being.
            if as_archive:
                # Fake compressor then
                @contextmanager
                def overall_compressor(f):
                    yield f

                file_compressor = compressor
                is_archive = True
            else:
                overall_compressor = compressor

                @contextmanager
                def file_compressor(f):
                    yield f

                is_archive = False

            with overall_compressor(output_file) as compressed_file:
                for count, filename_full in enumerate(file_list, start=1):
                    payload_size += _attachOnefilePayloadFile(
                        output_file=compressed_file,
                        is_archive=is_archive,
                        file_compressor=file_compressor,
                        is_compressing=compression_indicator == b"Y",
                        use_compression_cache=use_compression_cache,
                        low_memory=low_memory,
                        filename_full=filename_full,
                        count=count,
                        dist_dir=dist_dir,
                        filename_encoding=filename_encoding,
                        file_checksums=file_checksums,
                        win_path_sep=win_path_sep,
                    )

                # Using empty filename as a terminator.
                filename_encoded = "\0".encode(filename_encoding)
                compressed_file.write(filename_encoded)
                payload_size += len(filename_encoded)

            # TODO: If put into a resource, this is not really needed anymore.
            if isWin32Windows():
                # add padding to have the start position at a double world boundary
                # this is needed on windows so that a possible certificate immediately
                # follows the start position
                pad = output_file.tell() % 8
                if pad != 0:
                    output_file.write(bytes(8 - pad))

            output_file.seek(0, 2)
            end_pos = output_file.tell()

            compressed_size = end_pos - start_pos

            if compression_indicator == b"Y":
                onefile_logger.info(
                    "Onefile payload compression ratio (%.2f%%) size %d to %d."
                    % (
                        (float(compressed_size) / payload_size) * 100,
                        payload_size,
                        compressed_size,
                    )
                )

            # Size of the payload data plus the size of that size storage, so C code can
            # jump directly to it.
            output_file.write(struct.pack("Q", end_pos - start_pos))

        closeProgressBar()

    _attachOnefilePayload()


def main():
    # Internal tool, most simple command line handling. This is the build directory
    # where main Nuitka put the .const files.
    dist_dir = sys.argv[1]
    onefile_output_filename = sys.argv[2]
    start_binary = os.path.normpath(sys.argv[3])  # Might switch from MSYS2 to CPython
    include_start_binary = sys.argv[4] == "True"
    file_checksums = sys.argv[5] == "True"
    win_path_sep = sys.argv[6] == "True"
    low_memory = sys.argv[7] == "True"
    as_archive = sys.argv[8] == "True"
    use_compression_cache = sys.argv[9] == "True"
    job_limit = int(sys.argv[10])
    expected_files_filename = sys.argv[11]

    expected_files = tuple(loadJsonFromFilename(expected_files_filename))

    enableProgressBar(os.getenv("NUITKA_PROGRESS_BAR", "none"))

    attachOnefilePayload(
        dist_dir=dist_dir,
        onefile_output_filename=onefile_output_filename,
        start_binary=start_binary,
        include_start_binary=include_start_binary,
        # We wouldn't be here, if that was not the case.
        expect_compression=True,
        as_archive=as_archive,
        use_compression_cache=use_compression_cache,
        file_checksums=file_checksums,
        win_path_sep=win_path_sep,
        low_memory=low_memory,
        job_limit=job_limit,
        expected_files=expected_files,
    )

    sys.exit(0)


#     Part of "Nuitka", an optimizing Python compiler that is compatible and
#     integrates with CPython, but also works on its own.
#
#     Licensed under the GNU Affero General Public License, Version 3 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#        https://www.gnu.org/licenses/agpl-3.0.txt
#
#     See also: "Nuitka Runtime Library Exception, Version 1.0" in file
#     "LICENSE-RUNTIME.txt" for additional permissions granted under Section 7.
#
#     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.
