# Copyright 2018 The Bazel Authors. All rights reserved.
#
# 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.

from __future__ import annotations

import argparse
import base64
import csv
import hashlib
import io
import os
import re
import stat
import sys
import zipfile
from pathlib import Path

_ZIP_EPOCH = (1980, 1, 1, 0, 0, 0)


def commonpath(path1, path2):
    ret = []
    for a, b in zip(path1.split(os.path.sep), path2.split(os.path.sep)):
        if a != b:
            break
        ret.append(a)
    return os.path.sep.join(ret)


def escape_filename_segment(segment):
    """Escapes a filename segment per https://www.python.org/dev/peps/pep-0427/#escaping-and-unicode

    This is a legacy function, kept for backwards compatibility,
    and may be removed in the future. See `escape_filename_distribution_name`
    and `normalize_pep440` for the modern alternatives.
    """
    return re.sub(r"[^\w\d.]+", "_", segment, re.UNICODE)


def normalize_package_name(name):
    """Normalize a package name according to the Python Packaging User Guide.

    See https://packaging.python.org/en/latest/specifications/name-normalization/
    """
    return re.sub(r"[-_.]+", "-", name).lower()


def escape_filename_distribution_name(name):
    """Escape the distribution name component of a filename.

    See https://packaging.python.org/en/latest/specifications/binary-distribution-format/#escaping-and-unicode
    """
    return normalize_package_name(name).replace("-", "_")


def normalize_pep440(version):
    """Normalize version according to PEP 440, with fallback for placeholders.

    If there's a placeholder in braces, such as {BUILD_TIMESTAMP},
    replace it with 0. Such placeholders can be used with stamping, in
    which case they would have been resolved already by now; if they
    haven't, we're doing an unstamped build, but we still need to
    produce a valid version. If such replacements are made, the
    original version string, sanitized to dot-separated alphanumerics,
    is appended as a local version segment, so you understand what
    placeholder was involved.

    If that still doesn't produce a valid version, use version 0 and
    append the original version string, sanitized to dot-separated
    alphanumerics, as a local version segment.

    """

    import packaging.version

    try:
        return str(packaging.version.Version(version))
    except packaging.version.InvalidVersion:
        pass

    sanitized = re.sub(r"[^a-z0-9]+", ".", version.lower()).strip(".")
    substituted = re.sub(r"\{\w+\}", "0", version)
    delimiter = "." if "+" in substituted else "+"
    try:
        return str(packaging.version.Version(f"{substituted}{delimiter}{sanitized}"))
    except packaging.version.InvalidVersion:
        return str(packaging.version.Version(f"0+{sanitized}"))


def arcname_from(
    name: str,
    distribution_prefix: str,
    strip_path_prefixes: Sequence[str] = (),  # noqa: F821
    add_path_prefix: str = "",
) -> str:
    """Return the within-archive name for a given file path name.

    Prefixes to strip are checked in order and only the first match will be used.

    Args:
        name: The file path eg 'mylib/a/b/c/file.py'
        distribution_prefix: The
        strip_path_prefixes: Remove these prefixes from names.
        add_path_prefix: Add prefix after stripping the path from names.
    """
    # Always use unix path separators.
    normalized_arcname = name.replace(os.path.sep, "/")
    # Don't manipulate names filenames in the .distinfo or .data directories.
    if distribution_prefix and normalized_arcname.startswith(distribution_prefix):
        return normalized_arcname
    for prefix in strip_path_prefixes:
        if normalized_arcname.startswith(prefix):
            return add_path_prefix + normalized_arcname[len(prefix) :]

    return add_path_prefix + normalized_arcname


class _WhlFile(zipfile.ZipFile):
    def __init__(
        self,
        filename,
        *,
        mode,
        distribution_prefix: str,
        strip_path_prefixes=None,
        add_path_prefix=None,
        compression=zipfile.ZIP_DEFLATED,
        quote_all_filenames: bool = False,
        **kwargs,
    ):
        self._distribution_prefix = distribution_prefix

        self._strip_path_prefixes = strip_path_prefixes or []
        self._add_path_prefix = add_path_prefix or ""
        # Entries for the RECORD file as (filename, digest, size) tuples.
        self._record: list[tuple[str, str, str]] = []
        # Whether to quote filenames in the RECORD file (for compatibility with
        # some wheels like torch that have quoted filenames in their RECORD).
        self.quote_all_filenames = quote_all_filenames

        super().__init__(filename, mode=mode, compression=compression, **kwargs)

    def distinfo_path(self, basename):
        return f"{self._distribution_prefix}.dist-info/{basename}"

    def data_path(self, basename):
        return f"{self._distribution_prefix}.data/{basename}"

    def add_file(self, package_filename, real_filename):
        """Add given file to the distribution."""

        if os.path.isdir(real_filename):
            directory_contents = sorted(os.listdir(real_filename))
            for file_ in directory_contents:
                self.add_file(
                    "{}/{}".format(package_filename, file_),
                    "{}/{}".format(real_filename, file_),
                )
            return

        arcname = arcname_from(
            package_filename,
            distribution_prefix=self._distribution_prefix,
            strip_path_prefixes=self._strip_path_prefixes,
            add_path_prefix=self._add_path_prefix,
        )
        zinfo = self._zipinfo(arcname)

        # Write file to the zip archive while computing the hash and length
        hash = hashlib.sha256()
        size = 0
        with open(real_filename, "rb") as fsrc:
            with self.open(zinfo, "w", force_zip64=True) as fdst:
                while True:
                    block = fsrc.read(2**20)
                    if not block:
                        break
                    fdst.write(block)
                    hash.update(block)
                    size += len(block)

        self._add_to_record(arcname, self._serialize_digest(hash), size)

    def add_string(self, filename, contents):
        """Add given 'contents' as filename to the distribution."""
        if isinstance(contents, str):
            contents = contents.encode("utf-8", "surrogateescape")
        zinfo = self._zipinfo(filename)
        self.writestr(zinfo, contents)
        hash = hashlib.sha256()
        hash.update(contents)
        self._add_to_record(filename, self._serialize_digest(hash), len(contents))

    def _serialize_digest(self, hash) -> str:
        # https://www.python.org/dev/peps/pep-0376/#record
        # "base64.urlsafe_b64encode(digest) with trailing = removed"
        digest = base64.urlsafe_b64encode(hash.digest())
        digest = b"sha256=" + digest.rstrip(b"=")
        return digest.decode("utf-8", "surrogateescape")

    def _add_to_record(self, filename: str, hash: str, size: int) -> None:
        self._record.append((filename, hash, str(size)))

    def _zipinfo(self, filename):
        """Construct deterministic ZipInfo entry for a file named filename"""
        # Strip leading path separators to mirror ZipInfo.from_file behavior
        separators = os.path.sep
        if os.path.altsep is not None:
            separators += os.path.altsep
        arcname = filename.lstrip(separators)

        zinfo = zipfile.ZipInfo(filename=arcname, date_time=_ZIP_EPOCH)
        zinfo.create_system = 3  # ZipInfo entry created on a unix-y system
        # Both pip and installer expect the regular file bit to be set in order for the
        # executable bit to be preserved after extraction
        # https://github.com/pypa/pip/blob/23.3.2/src/pip/_internal/utils/unpacking.py#L96-L100
        # https://github.com/pypa/installer/blob/0.7.0/src/installer/sources.py#L310-L313
        zinfo.external_attr = (
            stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO | stat.S_IFREG
        ) << 16  # permissions: -rwxrwxrwx
        zinfo.compress_type = self.compression
        return zinfo

    def _quote_filename(self, filename: str) -> str:
        """Return a possibly quoted filename for RECORD file."""
        filename = filename.lstrip("/")
        # Some RECORDs like torch have *all* filenames quoted and we must minimize diff.
        # Otherwise, we quote only when necessary (e.g. for filenames with commas).
        quoting = csv.QUOTE_ALL if self.quote_all_filenames else csv.QUOTE_MINIMAL
        with io.StringIO() as buf:
            csv.writer(buf, quoting=quoting).writerow([filename])
            return buf.getvalue().strip()

    def add_recordfile(self) -> str:
        """Write RECORD file to the distribution."""
        record_path = self.distinfo_path("RECORD")
        entries = self._record + [(record_path, "", "")]
        entries = [
            (self._quote_filename(fname), digest, size)
            for fname, digest, size in entries
        ]
        contents = "\n".join(",".join(entry) for entry in entries) + "\n"
        self.add_string(record_path, contents)
        return contents


class WheelMaker(object):
    def __init__(
        self,
        name,
        version,
        build_tag,
        python_tag,
        abi,
        platform,
        compress,
        outfile=None,
        strip_path_prefixes=None,
        add_path_prefix=None,
    ):
        self._name = name
        self._version = normalize_pep440(version)
        self._build_tag = build_tag
        self._python_tag = python_tag
        self._abi = abi
        self._platform = platform
        self._outfile = outfile
        self._strip_path_prefixes = strip_path_prefixes
        self._add_path_prefix = add_path_prefix
        self._compress = compress
        self._wheelname_fragment_distribution_name = escape_filename_distribution_name(
            self._name
        )

        self._distribution_prefix = (
            self._wheelname_fragment_distribution_name + "-" + self._version
        )

        self._whlfile = None

    def __enter__(self):
        self._whlfile = _WhlFile(
            self.filename(),
            mode="w",
            distribution_prefix=self._distribution_prefix,
            strip_path_prefixes=self._strip_path_prefixes,
            add_path_prefix=self._add_path_prefix,
            compression=(
                zipfile.ZIP_DEFLATED if self._compress else zipfile.ZIP_STORED
            ),
        )
        return self

    def __exit__(self, type, value, traceback):
        self._whlfile.close()
        self._whlfile = None

    def wheelname(self) -> str:
        components = [
            self._wheelname_fragment_distribution_name,
            self._version,
        ]
        if self._build_tag:
            components.append(self._build_tag)
        components += [self._python_tag, self._abi, self._platform]
        return "-".join(components) + ".whl"

    def filename(self) -> str:
        if self._outfile:
            return self._outfile
        return self.wheelname()

    def disttags(self):
        return ["-".join([self._python_tag, self._abi, self._platform])]

    def distinfo_path(self, basename):
        return self._whlfile.distinfo_path(basename)

    def data_path(self, basename):
        return self._whlfile.data_path(basename)

    def add_file(self, package_filename, real_filename):
        """Add given file to the distribution."""
        self._whlfile.add_file(package_filename, real_filename)

    def add_wheelfile(self):
        """Write WHEEL file to the distribution"""
        # TODO(pstradomski): Support non-purelib wheels.
        wheel_contents = """\
Wheel-Version: 1.0
Generator: bazel-wheelmaker 1.0
Root-Is-Purelib: {}
""".format("true" if self._platform == "any" else "false")
        for tag in self.disttags():
            wheel_contents += "Tag: %s\n" % tag
        self._whlfile.add_string(self.distinfo_path("WHEEL"), wheel_contents)

    def add_metadata(self, metadata, name, description):
        """Write METADATA file to the distribution."""
        # https://www.python.org/dev/peps/pep-0566/
        # https://packaging.python.org/specifications/core-metadata/
        metadata = re.sub("^Name: .*$", "Name: %s" % name, metadata, flags=re.MULTILINE)
        metadata += "Version: %s\n\n" % self._version
        # setuptools seems to insert UNKNOWN as description when none is
        # provided.
        metadata += description if description else "UNKNOWN"
        metadata += "\n"
        self._whlfile.add_string(self.distinfo_path("METADATA"), metadata)

    def add_recordfile(self):
        """Write RECORD file to the distribution."""
        self._whlfile.add_recordfile()


def get_files_to_package(input_files):
    """Find files to be added to the distribution.

    input_files: list of pairs (package_path, real_path)
    """
    files = {}
    for package_path, real_path in input_files:
        files[package_path] = real_path
    return files


def get_new_requirement_line(reqs_text: str, extra: str) -> str:
    """Formats a requirement text into a Requires-Dist metadata line."""
    # This is not imported at the top of the file due to the reliance
    # on this file in the `whl_library` repository rule which does not
    # provide `packaging` but does import symbols defined here.
    from packaging.requirements import Requirement

    req = Requirement(reqs_text.strip())
    req_extra_deps = f"[{','.join(req.extras)}]" if req.extras else ""

    # Handle URL requirements (PEP 508)
    if req.url:
        req_spec = f" @ {req.url}"
    else:
        req_spec = str(req.specifier)

    base = f"Requires-Dist: {req.name}{req_extra_deps}{req_spec}"

    if req.marker:
        if extra:
            return f"{base}; ({req.marker}) and {extra}"
        else:
            return f"{base}; {req.marker}"
    elif extra:
        return f"{base}; {extra}"
    else:
        return base


def resolve_argument_stamp(
    argument: str, volatile_status_stamp: Path, stable_status_stamp: Path
) -> str:
    """Resolve workspace status stamps format strings found in the argument string

    Args:
        argument (str): The raw argument represenation for the wheel (may include stamp variables)
        volatile_status_stamp (Path): The path to a volatile workspace status file
        stable_status_stamp (Path): The path to a stable workspace status file

    Returns:
        str: A resolved argument string
    """
    lines = (
        volatile_status_stamp.read_text().splitlines()
        + stable_status_stamp.read_text().splitlines()
    )
    for line in lines:
        if not line:
            continue
        key, value = line.split(" ", maxsplit=1)
        stamp = "{" + key + "}"
        argument = argument.replace(stamp, value)

    return argument


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Builds a python wheel")
    metadata_group = parser.add_argument_group("Wheel name, version and platform")
    metadata_group.add_argument(
        "--name", required=True, type=str, help="Name of the distribution"
    )
    metadata_group.add_argument(
        "--version", required=True, type=str, help="Version of the distribution"
    )
    metadata_group.add_argument(
        "--build_tag",
        type=str,
        default="",
        help="Optional build tag for the distribution",
    )
    metadata_group.add_argument(
        "--python_tag",
        type=str,
        default="py3",
        help="Python version, e.g. 'py2' or 'py3'",
    )
    metadata_group.add_argument("--abi", type=str, default="none")
    metadata_group.add_argument(
        "--platform", type=str, default="any", help="Target platform. "
    )

    output_group = parser.add_argument_group("Output file location")
    output_group.add_argument(
        "--out", type=str, default=None, help="Override name of ouptut file"
    )
    output_group.add_argument(
        "--no_compress",
        action="store_true",
        help="Disable compression of the final archive",
    )
    output_group.add_argument(
        "--name_file",
        type=Path,
        help="A file where the canonical name of the wheel will be written",
    )

    output_group.add_argument(
        "--strip_path_prefix",
        type=str,
        action="append",
        default=[],
        help="Path prefix to be stripped from input package files' path. "
        "Can be supplied multiple times. Evaluated in order.",
    )
    output_group.add_argument(
        "--path_prefix",
        type=str,
        default="",
        help="Path prefix to be prepended to input package files' path. "
        "It is prepended after stripping any specified path prefixes first.",
    )

    wheel_group = parser.add_argument_group("Wheel metadata")
    wheel_group.add_argument(
        "--metadata_file",
        type=Path,
        help="Contents of the METADATA file (before appending contents of "
        "--description_file)",
    )
    wheel_group.add_argument(
        "--description_file", help="Path to the file with package description"
    )
    wheel_group.add_argument(
        "--description_content_type",
        help="Content type of the package description",
    )
    wheel_group.add_argument(
        "--entry_points_file",
        help="Path to a correctly-formatted entry_points.txt file",
    )

    contents_group = parser.add_argument_group("Wheel contents")
    contents_group.add_argument(
        "--input_file",
        action="append",
        help="'package_path;real_path' pairs listing "
        "files to be included in the wheel. "
        "Can be supplied multiple times.",
    )
    contents_group.add_argument(
        "--input_file_list",
        action="append",
        help="A file that has all the input files defined as a list to avoid "
        "the long command",
    )
    contents_group.add_argument(
        "--extra_distinfo_file",
        action="append",
        help="'filename;real_path' pairs listing extra files to include in"
        "dist-info directory. Can be supplied multiple times.",
    )
    contents_group.add_argument(
        "--data_files",
        action="append",
        help="'filename;real_path' pairs listing data files to include in"
        "data directory. Can be supplied multiple times.",
    )

    build_group = parser.add_argument_group("Building requirements")
    build_group.add_argument(
        "--volatile_status_file",
        type=Path,
        help="Pass in the stamp info file for stamping",
    )
    build_group.add_argument(
        "--stable_status_file",
        type=Path,
        help="Pass in the stamp info file for stamping",
    )

    return parser.parse_args(sys.argv[1:])


def _parse_file_pairs(content: List[str]) -> List[List[str]]:  # noqa: F821
    """
    Parse ; delimited lists of files into a 2D list.
    """
    return [i.split(";", maxsplit=1) for i in content or []]


def main() -> None:
    arguments = parse_args()

    input_files = _parse_file_pairs(arguments.input_file)
    extra_distinfo_file = _parse_file_pairs(arguments.extra_distinfo_file)
    data_files = _parse_file_pairs(arguments.data_files)

    for input_file in arguments.input_file_list:
        with open(input_file) as _file:
            input_file_list = _file.read().splitlines()
        for _input_file in input_file_list:
            input_files.append(_input_file.split(";"))

    all_files = get_files_to_package(input_files)
    # Sort the files for reproducible order in the archive.
    all_files = sorted(all_files.items())

    strip_prefixes = [p for p in arguments.strip_path_prefix]

    if arguments.volatile_status_file and arguments.stable_status_file:
        name = resolve_argument_stamp(
            arguments.name,
            arguments.volatile_status_file,
            arguments.stable_status_file,
        )
    else:
        name = arguments.name

    if arguments.volatile_status_file and arguments.stable_status_file:
        version = resolve_argument_stamp(
            arguments.version,
            arguments.volatile_status_file,
            arguments.stable_status_file,
        )
    else:
        version = arguments.version

    with WheelMaker(
        name=name,
        version=version,
        build_tag=arguments.build_tag,
        python_tag=arguments.python_tag,
        abi=arguments.abi,
        platform=arguments.platform,
        outfile=arguments.out,
        strip_path_prefixes=strip_prefixes,
        add_path_prefix=arguments.path_prefix,
        compress=not arguments.no_compress,
    ) as maker:
        for package_filename, real_filename in all_files:
            maker.add_file(package_filename, real_filename)
        maker.add_wheelfile()

        description = None
        if arguments.description_file:
            with open(
                arguments.description_file, "rt", encoding="utf-8"
            ) as description_file:
                description = description_file.read()

        metadata = arguments.metadata_file.read_text(encoding="utf-8")

        # Search for any `Requires-Dist` entries that refer to other files and
        # expand them.

        for meta_line in metadata.splitlines():
            if not meta_line.startswith("Requires-Dist: "):
                continue

            if not meta_line[len("Requires-Dist: ") :].startswith("@"):
                # This is a normal requirement.
                package, _, extra = meta_line[len("Requires-Dist: ") :].rpartition(";")
                if not package:
                    # This is when the package requirement does not have markers.
                    continue
                extra = extra.strip()
                metadata = metadata.replace(
                    meta_line, get_new_requirement_line(package, extra)
                )
                continue

            # This is a requirement that refers to a file.
            file, _, extra = meta_line[len("Requires-Dist: @") :].partition(";")
            extra = extra.strip()

            reqs = []
            for reqs_line in Path(file).read_text(encoding="utf-8").splitlines():
                reqs_text = reqs_line.strip()
                if not reqs_text or reqs_text.startswith(("#", "-")):
                    continue

                # Strip any comments
                reqs_text, _, _ = reqs_text.partition("#")

                reqs.append(get_new_requirement_line(reqs_text, extra))

            if reqs:
                metadata = metadata.replace(meta_line, "\n".join(reqs))
            # File is empty
            # So replace the meta_line entirely, including removing newline chars
            else:
                metadata = re.sub(
                    re.escape(meta_line) + r"(?:\r?\n)?", "", metadata, count=1
                )

        maker.add_metadata(
            metadata=metadata,
            name=name,
            description=description,
        )

        if arguments.entry_points_file:
            maker.add_file(
                maker.distinfo_path("entry_points.txt"),
                arguments.entry_points_file,
            )

        # Sort the files for reproducible order in the archive.
        for filename, real_path in sorted(data_files):
            maker.add_file(maker.data_path(filename), real_path)
        for filename, real_path in sorted(extra_distinfo_file):
            maker.add_file(maker.distinfo_path(filename), real_path)

        maker.add_recordfile()

        # Since stamping may otherwise change the target name of the
        # wheel, the canonical name (with stamps resolved) is written
        # to a file so consumers of the wheel can easily determine
        # the correct name.
        arguments.name_file.write_text(maker.wheelname())


if __name__ == "__main__":
    main()