#!/usr/bin/env python3

from __future__ import print_function
import shutil
import argparse
import sys
import os
import stat
import errno

from kobo.shortcuts import run
from productmd.composeinfo import ComposeInfo
from productmd.rpms import Rpms


class ComposeCheckError(Exception):
    """ "
    Raised when compose check fails.
    """

    pass


class ComposeCheck(object):
    """
    Checks the basic information about compose before promoting it.

    This is not real Compose CI, but rather basic sanity check.
    """

    def __init__(
        self,
        path,
        target,
        allow_unsigned=False,
        allow_finished_incomplete=False,
        compose_type="production",
    ):
        """
        Creates new ComposeCheck instance.

        :param str path: Path to Compose to check.
        :param str target: Target path where the promoted compose should be
            copied into.
        :param bool allow_unsigned: If True, compose with unsigned packages
            can be promoted.
        :param bool allow_finished_incomplete: If True, compose in FINISHED_INCOMPLETE
            state can be promoted.
        :param str compose_type: Compose with the type can be promoted.
        """
        self.path = path
        self.target = target
        self.allow_unsigned = allow_unsigned
        self.allow_finished_incomplete = allow_finished_incomplete
        self.compose_type = compose_type

    def check_status(self):
        """
        Raises ComposeCheckError if Compose STATUS is not FINISHED.
        """
        print("Checking compose STATUS.")

        allowed_statuses = ["FINISHED"]
        if self.allow_finished_incomplete:
            allowed_statuses.append("FINISHED_INCOMPLETE")

        status_path = os.path.join(self.path, "STATUS")
        with open(status_path, "r") as f:
            status = f.readline()[:-1]
            if status not in allowed_statuses:
                err_msg = "Compose is not in %s status." % (
                    " or ".join(allowed_statuses)
                )
                raise ComposeCheckError(err_msg)

    def check_compose_info(self):
        """
        Raises ComposeCheckError if Compose type is not "production".
        """
        print("Checking compose type.")
        ci = ComposeInfo()
        ci.load(os.path.join(self.path, "compose", "metadata", "composeinfo.json"))
        if ci.compose.type != self.compose_type:
            raise ComposeCheckError('Compose type is not "%s".' % self.compose_type)

    def check_rpms(self):
        """
        Raises ComposeCheckError if there are unsigned packages in the Compose.
        """
        if self.allow_unsigned:
            return

        print("Checking for unsigned RPMs.")
        rpms = Rpms()
        rpms.load(os.path.join(self.path, "compose", "metadata", "rpms.json"))
        for per_arch_rpms in rpms.rpms.values():
            for per_build_rpms in per_arch_rpms.values():
                for per_srpm_rpms in per_build_rpms.values():
                    for rpm in per_srpm_rpms.values():
                        if not rpm["sigkey"]:
                            err_msg = "Some RPMs are not signed."
                            raise ComposeCheckError(err_msg)

    def check_symlinks(self):
        """
        Raises ComposeCheckError if some symlink in the Compose cannot be resolved
        or if the symlink's target is not on the same device as Compose target
        directory.
        """
        print("Checking symlinks.")
        # The `self.target` can consist of multiple non-existing directories. We therefore
        # need to check parent directories until we hit existing directory. The last possible
        # path tried is "/" which should always exist.
        target_dirname = os.path.dirname(self.target)
        while not os.path.exists(target_dirname):
            target_dirname = os.path.dirname(target_dirname)

        for root, dirs, files in os.walk(self.path):
            for p in dirs + files:
                path = os.path.join(root, p)
                path_stat = os.lstat(path)
                if not stat.S_ISLNK(path_stat.st_mode):
                    continue

                real_path = os.readlink(path)
                abspath = os.path.normpath(
                    os.path.join(os.path.dirname(path), real_path)
                )
                try:
                    os.stat(abspath)
                except Exception as e:
                    err_msg = "Symlink cannot be resolved: %s: %s." % (path, e)
                    raise ComposeCheckError(err_msg)

    def run(self):
        """
        Runs the compose checks. Raises ComposeCheckError in case of failed check.
        """
        self.check_status()
        self.check_compose_info()
        self.check_rpms()
        self.check_symlinks()


class ComposePromotion(object):
    """
    Contains methods and data to promote compose.
    """

    def __init__(self, compose, target):
        """
        Creates new ComposePromotion instance.

        :param str compose: Path to Compose to promote.
        :param str target: Target path where the promoted compose should be
            copied into.
        """
        self.compose = compose
        self.target = target

        # Tuple in (symlink_path, hardlink_path) format:
        #  - symlink_path is full path to symlink in the `compose` tree.
        #  - hardlink_path is full path to new hardlink in the `target` tree.
        self.symlinks = []

    def _copytree_ignore(self, path, names):
        """
        Helper method for `shutil.copytree` to ignore symlinks when copying compose.

        This method also populates `self.symlinks`.
        """
        print("Copying files in %s." % path)
        ignored = []
        rel_path = os.path.relpath(path, self.compose)
        for name in names:
            file_path = os.path.join(path, name)
            if os.path.islink(file_path):
                ignored.append(name)
                hardlink_path = os.path.join(self.target, rel_path, name)
                self.symlinks.append((file_path, hardlink_path))
        return ignored

    def _replace_symlinks_with_hardlinks(self):
        """
        Copy symlinks from `compose` to `target` and replace them with hardlinks.
        """
        print("Replacing %d symlinks with hardlinks." % len(self.symlinks))
        for symlink, hardlink_path in self.symlinks:
            real_path = os.readlink(symlink)
            abspath = os.path.normpath(
                os.path.join(os.path.dirname(symlink), real_path)
            )
            try:
                os.link(abspath, hardlink_path)
            except OSError as ex:
                if ex.errno == errno.EXDEV:
                    shutil.copy2(abspath, hardlink_path)
                else:
                    raise

    def _run_hardlink(self, target):
        """Run hardlink on the final destination to save more space.

        This should help with images that are generated in work/ and hardlinked
        to compose/ during compose process.

        :param str target: the final destination of promoted compose.
        """
        # Make sure hardlink command is available and -x option is supported
        # (the default hardlink command in RHEL7 does not support -x option).
        hardlink = "/usr/sbin/hardlink"
        if not os.path.isfile(hardlink):
            return
        _, output = run([hardlink, "-h"], can_fail=True)
        if "-x" not in str(output):
            return

        cmd = [hardlink, "-c", "-vv", "-x", "^Packages$", target]
        run(cmd, stdout=True, show_cmd=True)

    def promote(self):
        """
        Promotes the compose.
        """
        shutil.copytree(args.compose, args.target, ignore=self._copytree_ignore)
        self._replace_symlinks_with_hardlinks()
        self._run_hardlink(args.target)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Promote ODCS compose.")
    parser.add_argument("compose", help="Path to compose to promote.")
    parser.add_argument("target", help="Path to target location")
    parser.add_argument(
        "--allow-unsigned", action="store_true", help="Allow unsigned RPMs."
    )
    parser.add_argument(
        "--allow-finished-incomplete",
        action="store_true",
        help="Allow compose in FINISHED_INCOMPLETE state.",
    )
    parser.add_argument(
        "--no-checks",
        action="store_true",
        help="WARN: Promote the compose without any checks.",
    )
    parser.add_argument(
        "--compose-type",
        default="production",
        help="Allowed compose type, default: production.",
    )
    args = parser.parse_args()

    args.compose = os.path.abspath(args.compose)
    args.target = os.path.abspath(args.target)

    if not args.no_checks:
        compose_check = ComposeCheck(
            args.compose,
            args.target,
            args.allow_unsigned,
            args.allow_finished_incomplete,
            args.compose_type,
        )
        try:
            compose_check.run()
        except ComposeCheckError as e:
            print("Compose validation error: %s" % str(e))
            sys.exit(1)

    print("Promoting compose")
    compose_promotion = ComposePromotion(args.compose, args.target)
    compose_promotion.promote()
