# wix.py - WiX installer functionality
#
# 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.

import collections
import json
import os
import pathlib
import re
import shutil
import subprocess
import typing
import uuid
import xml.dom.minidom

from .downloads import download_entry
from .py2exe import (
    build_py2exe,
    stage_install,
)
from .util import (
    extract_zip_to_directory,
    get_vc_environment,
    normalize_windows_version,
    process_install_rules,
    python_exe_info,
    read_version_py,
    sign_with_signtool,
    SourceDirs,
)


EXTRA_PACKAGES = {
    'dulwich',
    'distutils',
    'keyring',
    'pygments',
    'win32ctypes',
}

EXTRA_INCLUDES = {
    #'_curses',  # TODO: figure out why py3's py2exe doesn't see this
    '_curses_panel',
}

EXTRA_INSTALL_RULES = [
    ('{shellext_dir}/src/ThgShell*.dll', './'),
    ('{thg_dir}/contrib/windows/diff-scripts/*', 'diff-scripts/'),
    ('{thg_dir}/doc/build/chm/TortoiseHg.chm', 'doc/'),
    ('{thg_dir}/win32/thg-cmenu-*.reg', 'i18n/cmenu/'),
    # TODO: why in lib for these?
    # TODO: get rid of the naming inconsistency
    ('{winbuild_dir}/contrib/{kdiff3}.exe', 'lib/kdiff3.exe'),
    ('{winbuild_dir}/contrib/TortoisePlink-{arch}.exe', 'lib/TortoisePlink.exe'),
    ('{winbuild_dir}/contrib/Pageant-{arch}.exe', './Pageant.exe'),
    ('{winbuild_dir}/misc/hgbook.pdf', 'doc/'),
]

STAGING_REMOVE_FILES = [
]

SHORTCUTS = {
    # hg.1.html'
    'thg.file.c33f9b76_d19a_4d5a_8941_5896f056f10e': {
        'Name': 'Mercurial Command Reference',
    },
    # hgignore.5.html
    'thg.file.eca2fbb3_49b6_4e50_afe4_956de8ff1fd4': {
        'Name': 'Mercurial Ignore Files',
    },
    # hgrc.5.html
    'thg.file.787d254f_487e_4925_8fbe_f95d68b4d4f0': {
        'Name': 'Mercurial Configuration Files',
    },
}


def run_candle(wix, cwd, wxs, source_dir, defines=None):
    args = [
        str(wix / 'candle.exe'),
        '-nologo',
        str(wxs),
        '-dSourceDir=%s' % source_dir,
    ]

    if defines:
        args.extend('-d%s=%s' % define for define in sorted(defines.items()))

    subprocess.run(args, cwd=str(cwd), check=True)


def make_files_xml(staging_dir: pathlib.Path, is_x64) -> str:
    """Create XML string listing every file to be installed."""

    # We derive GUIDs from a deterministic file path identifier.
    # We shoehorn the name into something that looks like a URL because
    # the UUID namespaces are supposed to work that way (even though
    # the input data probably is never validated).

    doc = xml.dom.minidom.parseString(
        '<?xml version="1.0" encoding="utf-8"?>'
        '<Wix xmlns="http://schemas.microsoft.com/wix/2006/wi">'
        '</Wix>'
    )

    # Assemble the install layout by directory. This makes it easier to
    # emit XML, since each directory has separate entities.
    manifest = collections.defaultdict(dict)

    for root, dirs, files in os.walk(staging_dir):
        dirs.sort()

        root = pathlib.Path(root)
        rel_dir = root.relative_to(staging_dir)

        for i in range(len(rel_dir.parts)):
            parent = '/'.join(rel_dir.parts[0 : i + 1])
            manifest.setdefault(parent, {})

        for f in sorted(files):
            full = root / f
            manifest[str(rel_dir).replace('\\', '/')][full.name] = full

    component_groups = collections.defaultdict(list)

    # Now emit a <Fragment> for each directory.
    # Each directory is composed of a <DirectoryRef> pointing to its parent
    # and defines child <Directory>'s and a <Component> with all the files.
    for dir_name, entries in sorted(manifest.items()):
        # The directory id is derived from the path. But the root directory
        # is special.
        if dir_name == '.':
            parent_directory_id = 'INSTALLDIR'
        else:
            parent_directory_id = 'thg.dir.%s' % dir_name.replace(
                '/', '.'
            ).replace('-', '_')

        fragment = doc.createElement('Fragment')
        directory_ref = doc.createElement('DirectoryRef')
        directory_ref.setAttribute('Id', parent_directory_id)

        # Add <Directory> entries for immediate children directories.
        for possible_child in sorted(manifest.keys()):
            if (
                dir_name == '.'
                and '/' not in possible_child
                and possible_child != '.'
            ):
                child_directory_id = ('thg.dir.%s' % possible_child).replace(
                    '-', '_'
                )
                name = possible_child
            else:
                if not possible_child.startswith('%s/' % dir_name):
                    continue
                name = possible_child[len(dir_name) + 1 :]
                if '/' in name:
                    continue

                child_directory_id = 'thg.dir.%s' % possible_child.replace(
                    '/', '.'
                ).replace('-', '_')

            directory = doc.createElement('Directory')
            directory.setAttribute('Id', child_directory_id)
            directory.setAttribute('Name', name)
            directory_ref.appendChild(directory)

        # Add <Component>s for files in this directory.
        for rel, source_path in sorted(entries.items()):
            if dir_name == '.':
                full_rel = rel
            else:
                full_rel = '%s/%s' % (dir_name, rel)

            component_unique_id = (
                'https://www.tortoisehg.org/wix-installer/0/component/%s'
                % full_rel
            )
            component_guid = uuid.uuid5(uuid.NAMESPACE_URL, component_unique_id)
            component_id = 'thg.component.%s' % str(component_guid).replace(
                '-', '_'
            )

            component = doc.createElement('Component')

            component.setAttribute('Id', component_id)
            component.setAttribute('Guid', str(component_guid).upper())
            component.setAttribute('Win64', 'yes' if is_x64 else 'no')

            # Assign this component to a top-level group.
            if dir_name == '.':
                component_groups['ROOT'].append(component_id)
            elif '/' in dir_name:
                component_groups[dir_name[0 : dir_name.index('/')]].append(
                    component_id
                )
            else:
                component_groups[dir_name].append(component_id)

            unique_id = (
                'https://www.tortoisehg.org/wix-installer/0/%s' % full_rel
            )
            file_guid = uuid.uuid5(uuid.NAMESPACE_URL, unique_id)

            # IDs have length limits. So use GUID to derive them.
            file_guid_normalized = str(file_guid).replace('-', '_')
            file_id = 'thg.file.%s' % file_guid_normalized

            file_element = doc.createElement('File')
            file_element.setAttribute('Id', file_id)
            file_element.setAttribute('Source', str(source_path))
            file_element.setAttribute('KeyPath', 'yes')
            file_element.setAttribute('ReadOnly', 'yes')

            component.appendChild(file_element)
            directory_ref.appendChild(component)

        fragment.appendChild(directory_ref)
        doc.documentElement.appendChild(fragment)

    for group, component_ids in sorted(component_groups.items()):
        fragment = doc.createElement('Fragment')
        component_group = doc.createElement('ComponentGroup')
        component_group.setAttribute('Id', 'thg.group.%s' % group)

        for component_id in component_ids:
            component_ref = doc.createElement('ComponentRef')
            component_ref.setAttribute('Id', component_id)
            component_group.appendChild(component_ref)

        fragment.appendChild(component_group)
        doc.documentElement.appendChild(fragment)

    # Add <Shortcut> to files that have it defined.
    for file_id, metadata in sorted(SHORTCUTS.items()):
        els = doc.getElementsByTagName('File')
        els = [el for el in els if el.getAttribute('Id') == file_id]

        if not els:
            raise Exception('could not find File[Id=%s]' % file_id)

        for el in els:
            shortcut = doc.createElement('Shortcut')
            shortcut.setAttribute('Id', 'thg.shortcut.%s' % file_id)
            shortcut.setAttribute('Directory', 'ProgramMenuDir')
            shortcut.setAttribute('Icon', 'hgIcon.ico')
            shortcut.setAttribute('IconIndex', '0')
            shortcut.setAttribute('Advertise', 'yes')
            for k, v in sorted(metadata.items()):
                shortcut.setAttribute(k, v)

            el.appendChild(shortcut)

    return doc.toprettyxml()


def build_installer_py2exe(
    source_dirs: SourceDirs,
    python_exe: pathlib.Path,
    msi_name='tortoisehg',
    version=None,
    extra_packages_script=None,
    extra_wxs: typing.Optional[typing.Dict[str, str]] = None,
    extra_features: typing.Optional[typing.List[str]] = None,
    signing_info: typing.Optional[typing.Dict[str, str]] = None,
):
    """Build a WiX MSI installer using py2exe.

    ``source_dirs`` is the collection of paths to the TortoiseHg source tree and
    other included dependencies to use.
    ``arch`` is the target architecture. either ``x86`` or ``x64``.
    ``python_exe`` is the path to the Python executable to use/bundle.
    ``version`` is the TortoiseHg version string. If not defined,
    ``tortoisehg/util/__version__.py`` will be consulted.
    ``extra_packages_script`` is a command to be run to inject extra packages
    into the py2exe binary. It should stage packages into the virtualenv and
    print a null byte followed by a newline-separated list of packages that
    should be included in the exe.
    ``extra_wxs`` is a dict of {wxs_name: working_dir_for_wxs_build}.
    ``extra_features`` is a list of additional named Features to include in
    the build. These must match Feature names in one of the wxs scripts.
    """
    py_info = python_exe_info(python_exe)

    arch = 'x64' if py_info['arch'] == '64bit' else 'x86'

    thg_build_dir = source_dirs.original / 'build'

    requirements_txt = (
        source_dirs.original
        / 'contrib'
        / 'packaging'
        / 'requirements-windows-pyqt5-installer.txt'
    )

    build_py2exe(
        source_dirs,
        thg_build_dir,
        python_exe,
        'wix',
        requirements_txt,
        extra_packages=EXTRA_PACKAGES,
        extra_packages_script=extra_packages_script,
        extra_includes=EXTRA_INCLUDES,
    )

    # TODO: this is what winbuild uses, but Mercurial uses a predicatable uuid.
    productid = str(uuid.uuid4()).upper()

    build_shellext(source_dirs, sorted({'x86', arch}), productid, signing_info)

    build_dir = thg_build_dir / ('wix-%s' % arch)
    staging_dir = build_dir / 'stage'

    build_dir.mkdir(exist_ok=True)

    # Purge the staging directory for every build so packaging is pristine.
    if staging_dir.exists():
        print('purging %s' % staging_dir)
        shutil.rmtree(staging_dir)

    stage_install(source_dirs, staging_dir, lower_case=True)

    # We also install some extra files.
    dirs = {
        "arch": arch,
        "hg_dir": str(source_dirs.hg),
        "kdiff3": "kdiff3x64" if arch == "x64" else "kdiff3",
        "thg_dir": str(source_dirs.thg),
        "shellext_dir": str(source_dirs.shellext),
        "winbuild_dir": str(source_dirs.winbuild)
    }

    rules = [(src.format(**dirs), dest) for (src, dest) in EXTRA_INSTALL_RULES]

    # The source_dir doesn't matter much here because the paths are absolute.
    process_install_rules(rules, source_dirs.winbuild, staging_dir)

    # kdiff3 needs qt.conf to find platforms and imageformat plugins
    with open(staging_dir / 'lib' / 'qt.conf', 'w') as fp:
        fp.write('[Paths]\nPlugins = ..\n')

    # workaround for QTBUG 57687
    # TODO: See if this can be dropped once we're past 5.9 on Windows with py2
    shutil.copy(
        source_dirs.winbuild / "contrib" / "spawn.cmd",
        staging_dir / 'lib' / 'spawn.cmd'
    )

    # And remove some files we don't want.
    for f in STAGING_REMOVE_FILES:
        p = staging_dir / f
        if p.exists():
            print('removing %s' % p)
            p.unlink()

    wix_dir = source_dirs.thg / 'win32' / 'wix'

    wxs_entries = {
        wix_dir / 'diff-scripts.wxs': staging_dir / 'diff-scripts',
        wix_dir / 'dist-py3.wxs': staging_dir,
        wix_dir / 'icons.wxs': staging_dir / 'icons',
        wix_dir / 'thg-locale.wxs': staging_dir / 'locale',
        source_dirs.shellext / 'wix' / 'cmenu-i18n.wxs':
            staging_dir / 'i18n' / 'cmenu',

        # TODO: generate these in the staging area
        source_dirs.shellext / 'wix' / 'thgshell.wxs': source_dirs.shellext,

        wix_dir / 'help.wxs': staging_dir / 'helptext',
        wix_dir / 'templates.wxs': staging_dir / 'templates',
        wix_dir / 'locale.wxs': staging_dir / 'locale',
        wix_dir / 'i18n.wxs': staging_dir / 'i18n',
        wix_dir / 'doc.wxs': staging_dir / 'doc',
        wix_dir / 'contrib.wxs': staging_dir / 'contrib',
    }

    extra_wxs = dict(extra_wxs if extra_wxs else {})
    for wxs, source_dir in wxs_entries.items():
        extra_wxs[str(wxs)] = str(source_dir)

    return run_wix_packaging(
        source_dirs,
        build_dir,
        staging_dir,
        arch,
        version=version,
        productid=productid,
        msi_name=msi_name,
        extra_wxs=extra_wxs,
        extra_features=extra_features,
        signing_info=signing_info,
    )


def build_shellext(
    source_dirs: SourceDirs,
    archs: typing.Iterable[str],
    productid: str,
    signing_info: typing.Optional[typing.Dict[str, str]] = None,
):
    """Build the shell extension for the given architectures."""

    thgversion = read_version_py(source_dirs.thg)

    # TODO: sync the version generation in setup.py with hg's setup.py to
    #  generate modern style version strings to eliminate this slicing off of
    #  the extra info.  This cleans up something like 6.1.1+33-e6a340244a11.
    #  Since both the rc value (if present) commit count (if not) go into the
    #  4th component, we don't care about that because the 4th component isn't
    #  used here.
    thgversion = thgversion.split('+', 1)[0]

    # This may have 3 or 4 components
    nv = normalize_windows_version(thgversion).split('.')

    shellext_dir = source_dirs.shellext

    # Clear files added by previous installer runs
    subprocess.run(
        ["hg.exe", "--config", "extensions.purge=", "purge", "--all"],
        cwd=str(shellext_dir),
        check=True,
    )

    def compile(arch):
        print("*************** Compiling %s shellext ***************" % arch)

        customenv = {
            'DEBUG': '1',
            'THG_PLATFORM': arch,
            'THG_EXTRA_CPPFLAGS': '/DTHG_PRODUCT_ID=%s' % productid,
            'THG_EXTRA_RCFLAGS':
                '/dTHG_VERSION_FIRST=%s '
                '/dTHG_VERSION_SECOND=%s '
                '/dTHG_VERSION_THIRD=%s '
                '/dTHG_PRODUCT_ID="%s"' % (nv[0], nv[1], nv[2], productid),
            'CL': '/MP'  # Enable multiprocessor builds
        }

        env = get_vc_environment(arch == "x64")
        env.update(customenv)

        subprocess.run(
            ["nmake.exe", "/nologo", "/f", "Makefile.nmake", "clean"],
            shell=True,
            cwd=str(shellext_dir / 'src'),
            env=env,
            check=True,
        )
        subprocess.run(
            ["nmake.exe", "/nologo", "/f", "Makefile.nmake"],
            shell=True,
            cwd=str(shellext_dir / 'src'),
            env=env,
            check=True,
        )

        if signing_info:
            sign_with_signtool(
                shellext_dir / "src" / ("ThgShell%s.dll" % arch),
                "TortoiseHg %s %s shell extension" % (arch, thgversion),
                subject_name=signing_info["subject_name"],
                cert_path=signing_info["cert_path"],
                cert_password=signing_info["cert_password"],
                timestamp_url=signing_info["timestamp_url"],
            )

    for arch in archs:
        compile(arch)


def run_wix_packaging(
    source_dirs: SourceDirs,
    build_dir: pathlib.Path,
    staging_dir: pathlib.Path,
    arch: str,
    version: str,
    productid: str,
    msi_name: typing.Optional[str] = "tortoisehg",
    suffix: str = "",
    extra_wxs: typing.Optional[typing.Dict[str, str]] = None,
    extra_features: typing.Optional[typing.List[str]] = None,
    signing_info: typing.Optional[typing.Dict[str, str]] = None,
):
    """Invokes WiX to package up a built TortoiseHg.

    ``signing_info`` is a dict defining properties to facilitate signing the
    installer. Recognized keys include ``name``, ``subject_name``,
    ``cert_path``, ``cert_password``, and ``timestamp_url``. If populated,
    we will sign both the hg.exe and the .msi using the signing credentials
    specified.
    """

    orig_version = version or read_version_py(source_dirs.thg)
    # TODO: fix this hack that's also done in the shellext build
    orig_version = orig_version.split('+', 1)[0]
    version = normalize_windows_version(orig_version)
    print('using version string: %s' % version)
    if version != orig_version:
        print('(normalized from: %s)' % orig_version)

    if signing_info:
        # TODO: sign other executables
        sign_with_signtool(
            staging_dir / "hg.exe",
            "%s %s" % (signing_info["name"], version),
            subject_name=signing_info["subject_name"],
            cert_path=signing_info["cert_path"],
            cert_password=signing_info["cert_password"],
            timestamp_url=signing_info["timestamp_url"],
        )

    wix_dir = source_dirs.thg / 'win32' / 'wix'

    wix_pkg, wix_entry = download_entry('wix', build_dir)
    wix_path = build_dir / ('wix-%s' % wix_entry['version'])

    if not wix_path.exists():
        extract_zip_to_directory(wix_pkg, wix_path)

    # TODO: drop this path hacking when the following relative path errors are
    #  fixed (after py2 is dropped):
    #  dist-py3.wxs(106) : error LGHT0103 : The system cannot find the file '..\contrib\kdiff3x64.exe'.
    #  tortoisehg-py3.wxs(100) : error LGHT0103 : The system cannot find the file '..\extension-versions.txt'.
    #  tortoisehg-py3.wxs(116) : error LGHT0103 : The system cannot find the file '..\contrib\Pageant-x64.exe'.
    #  tortoisehg-py3.wxs(139) : error LGHT0103 : The system cannot find the file '..\misc\hgbook.pdf'.
    #  dist-py3.wxs(100) : error LGHT0103 : The system cannot find the file '..\contrib\kdiff3x64.exe'.
    #  dist-py3.wxs(110) : error LGHT0103 : The system cannot find the file '..\contrib\TortoisePlink-x64.exe'.
    #  thgshell.wxs(72) : error LGHT0103 : The system cannot find the file '../contrib/TortoiseOverlays\TortoiseOverlays-1.1.3.21564-x64.msm'.
    #  thgshell.wxs(76) : error LGHT0103 : The system cannot find the file '../contrib/TortoiseOverlays\TortoiseOverlays-1.1.3.21564-win32.msm'.
    build_dir = source_dirs.thg

    source_build_rel = pathlib.Path(os.path.relpath(source_dirs.thg, build_dir))

    defines = {'Platform': arch}

    # Derive a .wxs file with the staged files.
    # manifest_wxs = build_dir / 'stage.wxs'
    # with manifest_wxs.open('w', encoding='utf-8') as fh:
    #     fh.write(make_files_xml(staging_dir, is_x64=arch == 'x64'))
    #
    # run_candle(wix_path, build_dir, manifest_wxs, staging_dir, defines=defines)

    for source, rel_path in sorted((extra_wxs or {}).items()):
        run_candle(wix_path, build_dir, source, rel_path, defines=defines)

    source = wix_dir / 'tortoisehg-py3.wxs'

    # TODO: get the actual hg value
    hgversion = version

    defines['Version'] = version
    defines['Comments'] = 'Installs TortoiseHg %s, Mercurial %s on %s' % (
        version, hgversion, arch,
    )

    defines["PythonVersion"] = "3"
    defines['ProductId'] = productid
    defines['ShellextRepoFolder'] = str(source_dirs.shellext)

    # if (staging_dir / "lib").exists():
    #     defines["TortoiseHgHasLib"] = "1"

    if extra_features:
        assert all(';' not in f for f in extra_features)
        defines['TortoiseHgExtraFeatures'] = ';'.join(extra_features)

    run_candle(wix_path, build_dir, source, source_build_rel, defines=defines)

    msi_path = (
        source_dirs.original
        / 'dist'
        / ('%s-%s-%s%s.msi' % (msi_name, orig_version, arch, suffix))
    )

    args = [
        str(wix_path / 'light.exe'),
        '-nologo',
        '-ext',
        'WixUIExtension',
        '-sw1076',
        '-spdb',
        '-o',
        str(msi_path),
    ]

    for source, rel_path in sorted((extra_wxs or {}).items()):
        assert source.endswith('.wxs')
        source = os.path.basename(source)
        args.append(str(build_dir / ('%s.wixobj' % source[:-4])))

    args.extend(
        [
            # str(build_dir / 'stage.wixobj'),
            str(build_dir / 'tortoisehg-py3.wixobj'),
        ]
    )

    # TODO: drop this `.parent` path altering when relative paths are cleaned
    #  up in the *.wxs files
    subprocess.run(args, cwd=str(source_dirs.thg.parent), check=True)

    print('%s created' % msi_path)

    if signing_info:
        sign_with_signtool(
            msi_path,
            "%s %s" % (signing_info["name"], version),
            subject_name=signing_info["subject_name"],
            cert_path=signing_info["cert_path"],
            cert_password=signing_info["cert_password"],
            timestamp_url=signing_info["timestamp_url"],
        )

    return {
        'msi_path': msi_path,
    }
