# util.py - Common packaging utility code.
#
# Copyright 2019 Gregory Szorc <gregory.szorc@gmail.com>
# Copyright 2022 Matt Harbison <mharbison72@gmail.com>
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.

# no-check-code because Python 3 native.

from collections.abc import MutableMapping

import distutils.version
import getpass
import glob
import os
import pathlib
import re
import shutil
import subprocess
import tarfile
import zipfile


class SourceDirs:
    def __init__(self, source_dir: pathlib.Path):
        self._dependencies = source_dir / "dependencies"

        self.original = source_dir
        self.evolve = self._dependencies / "evolve"
        self.winbuild = self._dependencies / "winbuild"
        self.hg = self.winbuild / "build" / "hg"
        self.shellext = self.winbuild / "build" / "shellext"
        self.thg = self.winbuild / "build" / "thg"

    def clean(self):
        if self._dependencies.exists():
            print("Deleting dependencies")
            shutil.rmtree(self._dependencies)


def extract_tar_to_directory(source: pathlib.Path, dest: pathlib.Path):
    with tarfile.open(source, 'r') as tf:
        tf.extractall(dest)


def extract_zip_to_directory(source: pathlib.Path, dest: pathlib.Path):
    with zipfile.ZipFile(source, 'r') as zf:
        zf.extractall(dest)


def find_vc_installation() -> pathlib.Path:
    """Finds the latest Visual C++ installation."""
    # We invoke vswhere to find the latest Visual Studio install.
    vswhere = (
        pathlib.Path(os.environ["ProgramFiles(x86)"])
        / "Microsoft Visual Studio"
        / "Installer"
        / "vswhere.exe"
    )

    if not vswhere.exists():
        raise Exception(
            "could not find vswhere.exe: %s does not exist" % vswhere
        )

    args = [
        str(vswhere),
        # -products * is necessary to return results from Build Tools
        # (as opposed to full IDE installs).
        "-products",
        "*",
        "-requires",
        "Microsoft.VisualCpp.Redist.14.Latest",
        "-latest",
        "-property",
        "installationPath",
    ]

    return pathlib.Path(os.fsdecode(subprocess.check_output(args).strip()))


# Copied from os._Environ class, except the constructor.
class environ(MutableMapping):
    def __init__(self, data):
        # Where Env Var Names Must Be UPPERCASE
        def check_str(value):
            if not isinstance(value, str):
                raise TypeError("str expected, not %s" % type(value).__name__)
            return value

        def encodekey(key):
            return check_str(key).upper()

        self.encodekey = encodekey
        self.decodekey = str
        self.encodevalue = check_str
        self.decodevalue = str
        self._data = {}

        for key, value in data.items():
            self._data[encodekey(key)] = value

    def __getitem__(self, key):
        try:
            value = self._data[self.encodekey(key)]
        except KeyError:
            # raise KeyError with the original key value
            raise KeyError(key) from None
        return self.decodevalue(value)

    def __setitem__(self, key, value):
        key = self.encodekey(key)
        value = self.encodevalue(value)
        self._data[key] = value

    def __delitem__(self, key):
        encodedkey = self.encodekey(key)
        try:
            del self._data[encodedkey]
        except KeyError:
            # raise KeyError with the original key value
            raise KeyError(key) from None

    def __iter__(self):
        # list() from dict object is an atomic operation
        keys = list(self._data)
        for key in keys:
            yield self.decodekey(key)

    def __len__(self):
        return len(self._data)

    def __repr__(self):
        return 'environ({{{}}})'.format(', '.join(
            ('{!r}: {!r}'.format(self.decodekey(key), self.decodevalue(value))
            for key, value in self._data.items())))

    def copy(self):
        return environ(self._data)

    def setdefault(self, key, value):
        if key not in self:
            self[key] = value
        return self[key]


def get_vc_environment(x64=False) -> MutableMapping:
    """Fetch the build environment for the latest Visual C++ installation."""
    # This just gets us a path like
    # C:\Program Files (x86)\Microsoft Visual Studio\2019\Community
    # Actually vcvars*.bat is under a path like:
    # VC\Auxiliary\Build\vcvars<32|64>.bat.
    vs_install_path = find_vc_installation()

    vcvars_bat = (
        vs_install_path
        / "VC"
        / "Auxiliary"
        / "Build"
        / ("vcvars%s.bat" % ("64" if x64 else "32"))
    )

    output = subprocess.check_output(
        'cmd.exe /c ""%s" > NUL" && set' % vcvars_bat,
    ).strip()

    # Windows doesn't care about the case of environment variables, and some
    # key ones default to mixed case, like 'Path'.  Python does care because it
    # uses a `dict`, so wrap it in a simulation of `os._Environ` that internally
    # uppercases, but allows lookup using any case to let the rest of the code
    # access things consistently.
    #
    # https://bugs.python.org/issue28824
    env = {}
    for line in os.fsdecode(output).splitlines():
        try:
            k, v = line.split('=', 1)
            env[k] = v
        except ValueError:
            pass  # PS1, when run under MSYS, has a few newlines by default

    return environ(env)


def find_vc_runtime_dll(x64=False):
    """Finds Visual C++ Runtime DLL to include in distribution."""
    # This just gets us a path like
    # C:\Program Files (x86)\Microsoft Visual Studio\2019\Community
    # Actually vcruntime140.dll is under a path like:
    # VC\Redist\MSVC\<version>\<arch>\Microsoft.VC14<X>.CRT\vcruntime140.dll.
    vs_install_path = find_vc_installation()

    arch = "x64" if x64 else "x86"

    search_glob = (
        r"%s\VC\Redist\MSVC\*\%s\Microsoft.VC14*.CRT\vcruntime140.dll"
        % (vs_install_path, arch)
    )

    candidates = glob.glob(search_glob, recursive=True)

    for candidate in reversed(candidates):
        return pathlib.Path(candidate)

    raise Exception("could not find vcruntime140.dll")


PRINT_QT_BIN_DIR = """
import os, sys
join = os.path.join
print(join(sys.exec_prefix, 'lib', 'site-packages', 'PyQt5', 'Qt5', 'bin'))
""".strip()


def get_qt_dependencies(
    python_exe: pathlib.Path,
    dist_dir: pathlib.Path
):
    """Scan the Qt*/bin directory and enumerate the full path to dependencies of
    the bundled Qt native libraries that py2exe failed to stage in the named
    ``dist_dir``.  The set of libraries are architecture specific, and mostly
    related to the C runtime used to build Qt.
    """
    py_info = python_exe_info(python_exe)

    vc_x64 = py_info['arch'] == '64bit'

    env = get_vc_environment(x64=vc_x64)

    def dumpbin(file):
        dump = subprocess.check_output(
            ["dumpbin.exe", r'/IMPORTS', str(file)],
            shell=True,
            env=env,
        )

        for l in dump.splitlines():
            m = re.search(rb'^    ([^. ]+\.dll)', l, re.IGNORECASE)

            if not m:
                continue

            dll = os.fsdecode(m.group(1)).lower()

            if dll.startswith("api-"):  # Skip known low level DLLs
                continue

            yield dll

    seen_files = set()
    all_dependencies = set()

    # This casts a wider net than just Qt files, because not all imageformats/,
    # for example, start with "Qt".  It's not a big deal though because any
    # dependencies not found in the Qt/bin directory will be ignored.  And since
    # the PyQt5.*.pyd files import DLLs in Qt/bin that were already staged by
    # py2exe, track the already staged files so they are excluded from the end
    # result.
    for pattern in ["*.pyd", "*.dll"]:
        for file in dist_dir.rglob(pattern):
            seen_files.add(file.name.lower())
            all_dependencies.update(dll for dll in dumpbin(file))

    binpath = pathlib.Path(
        os.fsdecode(subprocess.check_output(
            [str(python_exe), '-c', PRINT_QT_BIN_DIR],
        ).strip())
    )

    # This must be getting called via LoadLibrary(), because none of *.dll or
    # *.pyd links against this.  It is required on 64-bit builds to connect
    # to the website and check for updates.  This will in turn pull in
    # libcrypto-1_1-x64.dll as a dependency.
    if vc_x64:
        all_dependencies.add('libssl-1_1-x64.dll')

    # seen_files were already processed, so they don't need to be rescanned.
    candidates = all_dependencies.difference(seen_files)
    dependencies = set()

    # The candidate list is updated on each iteration to handle dependencies of
    # dependencies.
    while len(candidates) > 0:
        for c in set(candidates):
            seen_files.add(c)

            abspath = binpath / c
            if abspath.exists():
                dependencies.add(abspath)
                all_dependencies.update(dll for dll in dumpbin(abspath))

        candidates = all_dependencies.difference(seen_files)

    return dependencies


def windows_10_sdk_info():
    """Resolves information about the Windows 10 SDK."""

    base = pathlib.Path(os.environ['ProgramFiles(x86)']) / 'Windows Kits' / '10'

    if not base.is_dir():
        raise Exception('unable to find Windows 10 SDK at %s' % base)

    # Find the latest version.
    bin_base = base / 'bin'

    versions = [v for v in os.listdir(bin_base) if v.startswith('10.')]
    version = sorted(versions, reverse=True)[0]

    bin_version = bin_base / version

    return {
        'root': base,
        'version': version,
        'bin_root': bin_version,
        'bin_x86': bin_version / 'x86',
        'bin_x64': bin_version / 'x64',
    }


def normalize_windows_version(version):
    """Normalize Mercurial version string so WiX/Inno accepts it.

    Version strings have to be numeric ``A.B.C[.D]`` to conform with MSI's
    requirements.

    We normalize RC version or the commit count to a 4th version component.
    We store this in the 4th component because ``A.B.C`` releases do occur
    and we want an e.g. ``5.3rc0`` version to be semantically less than a
    ``5.3.1rc2`` version. This requires always reserving the 3rd version
    component for the point release and the ``X.YrcN`` release is always
    point release 0.

    In the case of an RC and presence of ``+`` suffix data, we can't use both
    because the version format is limited to 4 components. We choose to use
    RC and throw away the commit count in the suffix. This means we could
    produce multiple installers with the same normalized version string.

    >>> normalize_windows_version("5.3")
    '5.3.0'

    >>> normalize_windows_version("5.3rc0")
    '5.3.0.0'

    >>> normalize_windows_version("5.3rc1")
    '5.3.0.1'

    >>> normalize_windows_version("5.3rc1+hg2.abcdef")
    '5.3.0.1'

    >>> normalize_windows_version("5.3+hg2.abcdef")
    '5.3.0.2'
    """
    if '+' in version:
        version, extra = version.split('+', 1)
    else:
        extra = None

    # 4.9rc0
    if version[:-1].endswith('rc'):
        rc = int(version[-1:])
        version = version[:-3]
    else:
        rc = None

    # Ensure we have at least X.Y version components.
    versions = [int(v) for v in version.split('.')]
    while len(versions) < 3:
        versions.append(0)

    if len(versions) < 4:
        if rc is not None:
            versions.append(rc)
        elif extra:
            # hg<commit count>.<hash>+<date>
            versions.append(int(extra.split('.')[0][2:]))

    return '.'.join('%d' % x for x in versions[0:4])


def find_signtool():
    """Find signtool.exe from the Windows SDK."""
    sdk = windows_10_sdk_info()

    for key in ('bin_x64', 'bin_x86'):
        p = sdk[key] / 'signtool.exe'

        if p.exists():
            return p

    raise Exception('could not find signtool.exe in Windows 10 SDK')


def sign_with_signtool(
    file_path,
    description,
    subject_name=None,
    cert_path=None,
    cert_password=None,
    timestamp_url=None,
):
    """Digitally sign a file with signtool.exe.

    ``file_path`` is file to sign.
    ``description`` is text that goes in the signature.

    The signing certificate can be specified by ``cert_path`` or
    ``subject_name``. These correspond to the ``/f`` and ``/n`` arguments
    to signtool.exe, respectively.

    The certificate password can be specified via ``cert_password``. If
    not provided, you will be prompted for the password.

    ``timestamp_url`` is the URL of a RFC 3161 timestamp server (``/tr``
    argument to signtool.exe).
    """
    if cert_path and subject_name:
        raise ValueError('cannot specify both cert_path and subject_name')

    while cert_path and not cert_password:
        cert_password = getpass.getpass('password for %s: ' % cert_path)

    args = [
        str(find_signtool()),
        'sign',
        '/v',
        '/fd',
        'sha256',
        '/d',
        description,
    ]

    if cert_path:
        args.extend(['/f', str(cert_path), '/p', cert_password])
    elif subject_name:
        args.extend(['/n', subject_name])

    if timestamp_url:
        args.extend(['/tr', timestamp_url, '/td', 'sha256'])

    args.append(str(file_path))

    print('signing %s' % file_path)
    subprocess.run(args, check=True)


PRINT_PYTHON_INFO = '''
import platform; print("%s:%s" % (platform.architecture()[0], platform.python_version()))
'''.strip()


def python_exe_info(python_exe: pathlib.Path):
    """Obtain information about a Python executable."""

    res = subprocess.check_output([str(python_exe), '-c', PRINT_PYTHON_INFO])

    arch, version = res.decode('utf-8').split(':')

    version = distutils.version.LooseVersion(version)

    return {
        'arch': arch,
        'version': version,
        'py3': version >= distutils.version.LooseVersion('3'),
    }


def process_install_rules(
    rules: list, source_dir: pathlib.Path, dest_dir: pathlib.Path
):
    for source, dest in rules:
        if '*' in source:
            if not dest.endswith('/'):
                raise ValueError('destination must end in / when globbing')

            # We strip off the source path component before the first glob
            # character to construct the relative install path.
            prefix_end_index = source[: source.index('*')].rindex('/')
            relative_prefix = source_dir / source[0:prefix_end_index]

            for res in glob.glob(str(source_dir / source), recursive=True):
                source_path = pathlib.Path(res)

                if source_path.is_dir():
                    continue

                rel_path = source_path.relative_to(relative_prefix)

                dest_path = dest_dir / dest[:-1] / rel_path

                dest_path.parent.mkdir(parents=True, exist_ok=True)
                print('copying %s to %s' % (source_path, dest_path))
                shutil.copy(source_path, dest_path)

        # Simple file case.
        else:
            source_path = pathlib.Path(source)

            if dest.endswith('/'):
                dest_path = pathlib.Path(dest) / source_path.name
            else:
                dest_path = pathlib.Path(dest)

            full_source_path = source_dir / source_path
            full_dest_path = dest_dir / dest_path

            full_dest_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(full_source_path, full_dest_path)
            print('copying %s to %s' % (full_source_path, full_dest_path))


def read_version_py(source_dir: pathlib.Path):
    """Read the tortoisehg/util/__version__.py file to resolve the version
    string.
    """
    p = source_dir / 'tortoisehg' / 'util' / '__version__.py'

    with p.open('r', encoding='utf-8') as fh:
        m = re.search('version = "([^"]+)"', fh.read(), re.MULTILINE)

        if not m:
            raise Exception('could not parse %s' % p)

        return m.group(1)
