#! /usr/bin/python3

"""
Take the JSON provided by Mock, download corresponding RPMs, and put them into
an RPM repository.
"""

# pylint: disable=invalid-name

import argparse
import concurrent.futures
import json
import logging
import os
import shutil
import subprocess
import sys

from urllib.parse import quote, urlparse
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry

import requests

logging.basicConfig(level=logging.DEBUG)
log = logging.getLogger(__name__)


def request_with_retry(retries=5, backoff_factor=0.3,
                       status_forcelist=(500, 502, 503, 504, 408, 429), session=None):
    """
    Wrapper for requests Session with default retries

    Converted from koji:
    https://pagure.io/koji/blob/fccf4fa3f990ca454a411c8a699e693f4e5b0ac8/f/koji/__init__.py#_2004
    originally stolen from
    https://www.peterbe.com/plog/best-practice-with-retries-with-requests
    """
    session = session or requests.Session()
    retry = Retry(total=retries, read=retries, connect=retries,
                  backoff_factor=backoff_factor,
                  status_forcelist=status_forcelist)
    adapter = HTTPAdapter(max_retries=retry)
    session.mount('http://', adapter)
    session.mount('https://', adapter)
    return session


def download_file(url, outputdir, client_auth_map):
    """
    Download a single file (pool worker)
    """
    basename = os.path.basename(url)
    file_name = os.path.join(outputdir, basename)
    dirname = os.path.dirname(url)

    # basename often contains '+' character, e.g., g++ package, url-encode it
    url = dirname + "/" + quote(basename)
    log.info("Downloading %s", url)

    try:
        parsed = urlparse(url)
        host_with_port = parsed.netloc
        host_only = parsed.hostname or host_with_port

        cert_arg = None
        # Prefer exact host:port mapping, then hostname-only
        mapping_value = client_auth_map.get(host_with_port) or client_auth_map.get(host_only)
        if mapping_value:
            cert_path, key_path = mapping_value
            cert_arg = (cert_path, key_path) if key_path else cert_path

        session = request_with_retry()
        request_kwargs = {"stream": True, "timeout": 60}
        if cert_arg is not None:
            request_kwargs["cert"] = cert_arg

        with session.get(url, **request_kwargs) as response:
            if response.status_code != 200:
                return False
            with open(file_name, "wb") as fd:
                shutil.copyfileobj(response.raw, fd)
            return True
    except:  # noqa: E722
        log.exception("Exception raised for %s", url)
        raise


def _argparser():
    parser = argparse.ArgumentParser(
        prog='mock-hermetic-repo',
        description=(
            "Prepare a repository for a `mock --hermetic-build` build. "
            "Given a Mock buildroot \"lockfile\"\n\n"
            "  a) create an output repo directory,\n"
            "  b) download and place all the necessary RPM files there,\n"
            "  c) create a local RPM repository there (run createrepo), and\n"
            "  d) dump there also the previously used bootstrap image as "
            "a tarball.\n\n"
            "Lockfile is a buildroot_lock.json file from Mock's "
            "result directory; it is a JSON file generated by the "
            "--calculate-build-dependencies option/buildroot_lock "
            "plugin."),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument("--lockfile", required=True,
                        help=(
                            "Select buildroot_lock.json filename on your system, "
                            "typically located in the Mock's result directory "
                            "upon the --calculate-build-dependencies mode "
                            "execution."))
    parser.add_argument("--output-repo", required=True,
                        help=(
                            "Download RPMs into this directory, and then run "
                            "/bin/createrepo_c utility there to populate the "
                            "RPM repo metadata."))
    parser.add_argument("--client-cert-for", dest="client_cert_for",
                        action="append", nargs='+', metavar="HOST CERT [KEY]",
                        help=(
                            "Register a client certificate (and optional key) to use for a "
                            "specific host. Repeatable. Usage: \n"
                            "  --client-cert-for HOST CERT [KEY]"))
    return parser


def prepare_image(image_specification, bootstrap_data, outputdir):
    """
    Store the tarball into the same directory where the RPMs are
    """
    # --scoped - it is better to refer will full repo name instead of just tag@sha
    #            also it matches with full pullspec from config's bootstrap_image
    # --preserve-digests - it is not needed for most images, but some could end with
    #            modified digests (semantically identical copy, but modified digest,
    #            e.g. embedded docker reference for v1 images), so better safe than sorry
    pull_cmd = ["skopeo", "sync", "--src", "docker", "--dest", "dir",
                "--scoped", "--preserve-digests"]
    # skopeo (currently) can't live with tag + sha, so strip the tag
    image_specification = image_specification.split(':')[0]
    if "architecture" in bootstrap_data:
        pull_cmd += ["--override-arch", bootstrap_data["architecture"]]
    if "pull_digest" in bootstrap_data:
        image_specification +=  "@" + bootstrap_data["pull_digest"]
    pull_cmd += [image_specification, outputdir]
    log.info("Pulling like: %s", ' '.join(pull_cmd))
    subprocess.check_output(pull_cmd)


def _main():
    options = _argparser().parse_args()

    # Build host -> (cert_path, key_path|None) mapping
    client_auth_map = {}
    if getattr(options, "client_cert_for", None):
        for parts in options.client_cert_for:
            if len(parts) < 2 or len(parts) > 3:
                log.error("Invalid --client-cert-for usage (expected: --client-cert-for HOST CERT [KEY]) -> %s", parts)
                sys.exit(2)
            host = parts[0]
            cert_path = parts[1]
            key_path = parts[2] if len(parts) == 3 else None
            if not cert_path:
                log.error("Invalid --client-cert-for value for %s (missing CERT path)", host)
                sys.exit(2)
            if not os.path.isfile(cert_path):
                log.error("Client certificate file does not exist for %s: %s", host, cert_path)
                sys.exit(2)
            if key_path and not os.path.isfile(key_path):
                log.error("Client key file does not exist for %s: %s", host, key_path)
                sys.exit(2)
            # Log registration without printing key path
            if key_path:
                log.info("Registering client certificate for %s: cert=%s (key provided)", host, cert_path)
            else:
                log.info("Registering client certificate for %s: cert=%s", host, cert_path)
            client_auth_map[host] = (cert_path, key_path)

    with open(options.lockfile, "r", encoding="utf-8") as fd:
        data = json.load(fd)

    try:
        os.makedirs(options.output_repo)
    except FileExistsError:
        pass

    failed = False
    urls = [i["url"] for i in data["buildroot"]["rpms"]]
    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        for i, out in zip(urls, executor.map(download_file,
                                             urls,
                                             [options.output_repo for _ in urls],
                                             [client_auth_map for _ in urls])):
            if out is False:
                log.error("Download failed: %s", i)
                failed = True
    if failed:
        log.error("RPM deps downloading failed")
        sys.exit(1)

    subprocess.check_call(["createrepo_c", "--compatibility", options.output_repo])

    prepare_image(data["config"]["bootstrap_image"], data["bootstrap"],
                  options.output_repo)


if __name__ == "__main__":
    _main()
