from Atomic import Atomic
from Atomic import util
from Atomic import info
from Atomic import update
from Atomic import verify
from Atomic import help as Help
from Atomic.mount import Mount
from Atomic.delete import Delete
import os
import math
import shutil
import tempfile
import argparse
from Atomic import backendutils

ATOMIC_CONFIG = util.get_atomic_config()
storage = ATOMIC_CONFIG.get('default_storage', "docker")

def convert_size(size):
    if size > 0:
        size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
        i = int(math.floor(math.log(size, 1000)))
        p = math.pow(1000, i)
        s = round(size/p, 2) # pylint: disable=round-builtin,old-division
        if s > 0:
            return '%s %s' % (s, size_name[i])
    return '0B'

def cli(subparser):
    # atomic images
    imagesp = subparser.add_parser("images",
                                   help=_("operate on images"))
    images_subparser = imagesp.add_subparsers(title='images subcommands',
                                              description="operate on images",
                                              help='additional help')

    # atomic images delete
    delete_parser = images_subparser.add_parser("delete",
                                                help=_("mark image for deletion"),
                                                epilog="Marks image. registry garbage-collection "
                                                "when invoked will recover used disk space")

    deletegroup = delete_parser.add_mutually_exclusive_group()
    deletegroup.add_argument("-f", "--force", default=False, dest="force",
                             action="store_true",
                             help=_("Force removal of local images, even if containers based on it exist.  Default is False"))

    deletegroup.add_argument("--remote", default=False, dest="remote",
                             action="store_true",
                             help=_("Delete image from remote repository"))

    delete_parser.add_argument("--storage", default=None, dest="storage",
                               help=_("Specify the storage from which to delete the image from. "
                                      "If not specified and there are images with the same name in "
                                      "different storages, you will be prompted to specify."))
    delete_parser.add_argument("-a", "--all", action='store_true',dest="all",
                               default=False, help=_("Delete all images"))
    delete_parser.add_argument("delete_targets", nargs='*',
                                    help=_("container image(s)"))
    delete_parser.set_defaults(_class=Delete, func='delete_image')

    if util.gomtree_available():
        generate_parser = images_subparser.add_parser("generate",
                                                  help=_("generate image 'manifests' if missing"))
        generate_parser.set_defaults(_class=Images, func='generate_validation_manifest')

        # Suppress storage flag: only docker images are supported for generate
        generate_parser.add_argument("--storage", default=storage, dest="storage",
                                     help=argparse.SUPPRESS)

    Help.cli(images_subparser)

    info.cli(images_subparser)

    # atomic images list
    list_parser = images_subparser.add_parser("list",
                                              help=_("list container images on your system"),
                                              epilog="atomic images by default will list all installed "
                                                     "container images on your system.")
    list_parser.set_defaults(_class=Images, func='display_all_image_info')

    list_parser.add_argument("-a", "--all", dest="all", default=False,
                             action="store_true",
                             help=_("Show all images, including intermediate images"))

    list_parser.add_argument("-f", "--filter", dest="filter", metavar="FILTER",
                             action="append",
                             help=_("Filter output based on VARIABLE=VALUE format"))

    list_parser.add_argument("-n", "--noheading", dest="heading", default=True,
                             action="store_false",
                             help=_("do not print heading when listing the images"))

    list_parser.add_argument("--no-trunc", dest="truncate", default=True,
                             action="store_false",
                             help=_("Don't truncate output"))

    list_parser.add_argument("-q", "--quiet", dest="quiet", default=False,
                             action="store_true",
                             help=_("Only display image IDs"))
    list_parser.add_argument("--json", action='store_true',dest="json", default=False,
                             help=_("print in a machine parseable form"))

    prune_parser = images_subparser.add_parser("prune",
                                               help=_("delete unused 'dangling' images"),
                                               epilog="Using the prune command, "
                                                      "will free up disk space deleting unused "
                                                      "'dangling' images")
    prune_parser.set_defaults(_class=Delete, func='prune_images')

    update.cli(images_subparser)

    verify.cli(images_subparser)

    info.cli_version(images_subparser)

class Images(Atomic):
    def __init__(self):
        super(Images, self).__init__()
        self.be_utils = backendutils.BackendUtils()

    def display_all_image_info(self):
        def get_col_lengths(_images):
            '''
            Determine the max length of the repository and tag names
            :param _images:
            :return: a set with len of repository and tag
            If there are no images, return 1, 1
            '''
            repo_tags = [y for x in _images if x.repotags for y in x.split_repotags]

            if repo_tags:
                return max([len(x[0]) for x in repo_tags]) + 2,\
                       max([len(x[1]) for x in repo_tags]) + 2
            else:
                return 1, 1

        if self.args.debug:
            util.write_out(str(self.args))
            self.be_utils.dump_backends()

        _images = self._get_images()
        for i in _images:
            i.repo, i.tag = i.split_repotags[0]

        if self.args.filter:
            _images = [x for x in _images if self._filter_include_image(x)]

        if self.args.json:
            util.output_json(self.return_json(_images))
            return 0

        if len(_images) == 0:
            return

        _max_repo, _max_tag = get_col_lengths(_images)

        if self.args.truncate:
            _max_id = 14
        else:
            _max_id = 65
        col_out = "{0:2} {1:" + str(_max_repo) + "} {2:" + str(_max_tag) + \
                  "} {3:" + str(_max_id) + "} {4:18} {5:14} {6:10}"

        if self.args.heading and not self.args.quiet:
            util.write_out(col_out.format(" ",
                                          "REPOSITORY",
                                          "TAG",
                                          "IMAGE ID",
                                          "CREATED",
                                          "VIRTUAL SIZE",
                                          "TYPE"))
        for image in _images:
            if self.args.quiet:
                util.write_out(image.id)

            else:
                indicator = ""
                if image.is_dangling:
                    indicator += "*"
                elif image.used:
                    indicator += ">"
                if image.vulnerable:
                    space = " " if len(indicator) < 1 else ""
                    if util.is_python2:
                        indicator = indicator + self.skull + space
                    else:
                        indicator = indicator + str(self.skull, "utf-8") + space
                _id = image.short_id if self.args.truncate else image.id
                util.write_out(col_out.format(indicator, image.repo or "<none>", image.tag or "<none>", _id, image.created[0:16],
                                              image.virtual_size, image.backend.backend))
        util.write_out("")
        return

    def _get_images(self):
        _images = self.be_utils.get_images(get_all=self.args.all)

        self._mark_used(_images)
        self._mark_vulnerable(_images)

        return _images

    def images(self):
        return self.return_json(self._get_images())


    def generate_validation_manifest(self):
        """
        Generates a gomtree validation manifest for a non-system image and stores it in
        ATOMIC_VAR_LIB
        :param:
        :return: None
        """
        _images = self.get_images(get_all=True)
        for image in _images:
            atomic_var_lib = util.ATOMIC_VAR_LIB
            if not image["RepoTags"]:
                continue
            iid = image["RepoTags"][0]
            if image["ImageType"] == "system":
                continue
            if iid == "<none>:<none>" or iid == "<none>":
                continue
            if os.path.exists(os.path.join(atomic_var_lib, "gomtree-manifests/%s.mtree" % iid)):
                continue
            manifestname = os.path.join(atomic_var_lib, "gomtree-manifests/%s.mtree" % iid)
            dname = os.path.dirname(manifestname)
            if not os.path.exists(dname):
                os.makedirs(dname)
            tmpdir = tempfile.mkdtemp()
            m = Mount()
            m.args = []
            m.image = iid
            m.mountpoint = tmpdir
            m.storage = self.args.storage
            m.mount()
            r = util.generate_validation_manifest(img_rootfs=tmpdir, keywords="type,uid,gid,mode,size,sha256digest")
            m.unmount()
            with open(manifestname,"wb",0) as f:
                f.write(r.stdout)
            shutil.rmtree(tmpdir)

    def _filter_include_image(self, image_obj):
        filterables = ["repo", "tag", "id", "image", "created", "size", "type", "is_dangling"]
        for i in self.args.filter:
            try:
                var, value = str(i).split("=")
            except ValueError:
                raise ValueError("The filter {} is not formatted correctly.  It should be VAR=VALUE".format(i))
            var = var.lower()
            if var == "repository":
                var = "repo"

            if var == "image":
                var = "id"

            if var not in filterables:
                raise ValueError("The filter {} is not valid.  "
                                 "Please choose from {}".format(var, filterables))
            if var == "type":
                var = "str_backend"

            if hasattr(image_obj, var) and value.lower() in str(getattr(image_obj, var)).lower():
                return True

        return False

    def _mark_used(self, images):
        assert isinstance(images, list)
        all_containers = [x.id for x in self.be_utils.get_containers()]
        for image in images:
            if image.id in all_containers:
                image.used = True

    def _mark_vulnerable(self, images):
        assert isinstance(images, list)
        vulnerable_uuids = self.get_vulnerable_ids()
        for image in images:
            if image.id in vulnerable_uuids:
                image.vulnerable = True

    def return_json(self, images):
        all_image_info = []
        all_vuln_info = self.get_all_vulnerable_info()
        keys = ['is_dangling', 'used', 'vulnerable', 'id', 'type', 'created', 'virtual_size']
        for img_obj in images:
            if not img_obj.repotags:
                continue
            img_dict = dict()
            img_dict['repo'], img_dict['tag'] = img_obj.repo, img_obj.tag
            for key in keys:
                img_dict[key] = getattr(img_obj, key, None)
            img_dict['vuln_info'] = \
                dict() if not img_obj.vulnerable else all_vuln_info.get(img_obj.id, None) # pylint: disable=no-member
            all_image_info.append(img_dict)
        return all_image_info


