#!/usr/bin/python3

import argparse
import ast
import collections.abc
import inspect
import os.path
import sys

ch_lib = "/usr/lib/charliecloud"
sys.path.insert(0, ch_lib)
import charliecloud as ch
import build
import build_cache as bu
import filesystem as fs
import image as im
import misc
import modify
import pull
import push


## Constants ##

# FIXME: It’s currently easy to get the ch-run path from another script, but
# hard from something in lib. So, we set it here for now.
ch.CH_BIN = os.path.dirname(os.path.abspath(
                 inspect.getframeinfo(inspect.currentframe()).filename))
ch.CH_RUN = ch.CH_BIN + "/ch-run"


## Main ##

def main():

   if (not os.path.exists(ch.CH_RUN)):
      ch.depfails.append(("missing", ch.CH_RUN))

   ap = ch.ArgumentParser(
      description="Build and manage images; completely unprivileged.",
      epilog="""Storage directory is used for caching and temporary images.
                Location: first defined of --storage, $CH_IMAGE_STORAGE, and
                %s.""" % fs.Storage.root_default(),
      sub_title="subcommands",
      sub_metavar="CMD")

   # Common options.
   #
   # --dependencies (and --help and --version) are options rather than
   # subcommands for consistency with other commands.
   #
   # These are also accepted *after* the subcommand, as it makes wrapping
   # ch-image easier and possibly improve the UX. There are multiple ways to
   # do this, though no tidy ones unfortunately. Here, we build up a
   # dictionary of options we want, and pass it to both main and subcommand
   # parsers; this works because both go into the same Namespace object. There
   # are two quirks to be aware of:
   #
   #   1. We omit the common options from subcommand --help for clarity and
   #      because before the subcommand is preferred.
   #
   #   2. We suppress defaults in the subcommand [1]. Without this, the
   #      subcommand option value wins even it it’s the default. :P Currently,
   #      if specified in both places, the subcommand value wins and the
   #      before value is not considered at all, e.g. "ch-image -vv foo -v"
   #      gives verbosity 1, not 3. This oddity seemed acceptable.
   #
   # Alternate approaches include:
   #
   #   * Set the main parser as the “parent” of the subcommand parser [2].
   #     This may be the documented approach? However, it adds all the
   #     subcommands to the subparser, which we don’t want. A workaround would
   #     be to create a *third* parser that’s the parent of both the main and
   #     subcommand parsers, but that seems like too much indirection to me.
   #
   #   * A two-stage parse (parse_known_args(), then parse_args() to have the
   #     main parser look again) works [3], but is complicated and has some
   #     odd side effects e.g. multiple subcommands will be accepted.
   #
   # Each sub-list is a group of options. They key identifies the mutually
   # exclustive group, or non-mutually exclusive if None.
   #
   # [1]: https://bugs.python.org/issue9351#msg373665
   # [2]: https://docs.python.org/3/library/argparse.html#parents
   # [3]: https://stackoverflow.com/a/54936198
   common_opts = \
      { ("bucache", "build cache common options"): [
           [["--cache"],
            { "action": "store_const",
              "const": ch.Build_Mode.ENABLED,
              "dest": "bucache",
              "help": "enable build cache" }],
           [["--no-cache"],
            { "action": "store_const",
              "const": ch.Build_Mode.DISABLED,
              "dest": "bucache",
              "help": "disable build cache" }],
           [["--rebuild"],
            { "action": "store_const",
              "const": ch.Build_Mode.REBUILD,
              "dest": "bucache",
              "help": "force cache misses for non-FROM instructions" }] ],
        (None, "misc common options"): [
           [["-a", "--arch"],
            { "metavar": "ARCH",
              "default": "host",
              "help": "architecture for image registries (default: host)"}],
           [["--always-download"],
            { "action": "store_true",
              "help": "redownload any image files when pulling"}],
           [["--auth"],
            { "action": "store_true",
              "help": "authenticated registry access; implied by push" }],
           [["--break"],
            { "metavar": "MODULE:LINE",
              "help": "break into PDB before LINE of MODULE" }],
           [["--cache-large"],
            { "metavar": "SIZE",
              "type": lambda s: ch.positive(s) * 2**20,  # internal unit: bytes
              "default": ch.positive(
                 os.environ.get("CH_IMAGE_CACHE_LARGE", 0)) * 2**20,
              "help": "large file threshold in MiB" }],
           [["--debug"],
            { "action": "store_true",
              "help": "add short traceback to fatal error hints" }],
           [["--dependencies"],
            { "action": misc.Dependencies,
              "help": "print any missing dependencies and exit" }],
           [["--no-lock"],
            { "action": "store_true",
              "help": "allow concurrent storage directory access (risky!)" }],
           [["--no-xattrs"],
            { "action": "store_true",
              "help": "disable xattrs and ACLs (overrides $CH_XATTRS)" }],
           [["--password-many"],
            { "action": "store_true",
              "help": "re-prompt each time a registry password is needed" }],
           [["--profile"],
            { "action": "store_true",
              "help": "dump profile to “./profile.{p,txt}”" }],
           [["-q", "--quiet"],
            { "action": "count",
              "default": 0,
              "help": "print less output (can be repeated)"}],
           [["-s", "--storage"],
            { "metavar": "DIR",
              "type": fs.Path,
              "help": "set builder internal storage directory to DIR" }],
           [["--tls-no-verify"],
            { "action": "store_true",
              "help": "don’t verify registry certificates (dangerous!)" }],
           [["-v", "--verbose"],
            { "action": "count",
              "default": 0,
              "help": "print extra chatter (can be repeated)" }],
           [["--version"],
            { "action": misc.Version,
              "help": "print version and exit" }],
           [["--xattrs"],
            { "action": "store_true",
              "help": "enable build cache support for xattrs and ACLs"}] ] }

   # Most, but not all, subcommands need to check dependencies before doing
   # anything (the exceptions being basic information commands like
   # storage-path). Similarly, only some need to initialize the storage
   # directory. These dictionaries map the dispatch function to a boolean
   # value saying whether to do those things.
   dependencies_check = dict()
   storage_init = dict()

   # Helper function to set up a subparser. The star forces the latter two
   # arguments to be called by keyword, for clarity.
   def add_opts(p, dispatch, *, deps_check, stog_init, help_=False):
      assert (not stog_init or deps_check)  # can’t init storage w/o deps
      if (dispatch is not None):
         p.set_defaults(func=dispatch)
         dependencies_check[dispatch] = deps_check
         storage_init[dispatch] = stog_init
      for ((name, title), group) in common_opts.items():
         if (name is None):
            p2 = p.add_argument_group(title=title)
         else:
            p2 = p.add_argument_group(title=title)
            p2 = p2.add_mutually_exclusive_group()
         for (args, kwargs) in group:
            if (help_):
               kwargs2 = kwargs
            else:
               kwargs2 = { **kwargs, "default": argparse.SUPPRESS }
            p2.add_argument(*args, **kwargs2)

   # main parser
   add_opts(ap, None, deps_check=False, stog_init=False, help_=True)

   # build
   sp = ap.add_parser("build", "build image from Dockerfile")
   add_opts(sp, build.main, deps_check=True, stog_init=True)
   sp.add_argument("-b", "--bind", metavar="SRC[:DST]",
                   action="append", default=[],
                   help="mount SRC at guest DST (default: same as SRC)")
   sp.add_argument("--build-arg", metavar="ARG[=VAL]",
                   action="append", default=[],
                   help="set build-time variable ARG to VAL, or $ARG if no VAL")
   sp.add_argument("-f", "--file", metavar="DOCKERFILE",
                   help="Dockerfile to use (default: CONTEXT/Dockerfile)")
   sp.add_argument("--force", metavar="MODE", nargs="?", default="seccomp",
                   type=ch.Force_Mode, const="seccomp",
                   help="inject unprivileged build workarounds")
   sp.add_argument("--force-cmd", metavar="CMD,ARG1[,ARG2...]",
                   action="append", default=[],
                   help="command arg(s) to add under --force=seccomp")
   sp.add_argument("-n", "--dry-run", action="store_true",
                   help="don’t execute instructions")
   sp.add_argument("--parse-only", action="store_true",
                   help="stop after parsing the Dockerfile")
   sp.add_argument("-t", "--tag", metavar="TAG",
                   help="name (tag) of image to create (default: inferred)")
   sp.add_argument("context", metavar="CONTEXT",
                   help="context directory")

   # build-cache
   sp = ap.add_parser("build-cache", "print build cache information")
   add_opts(sp, misc.build_cache, deps_check=True, stog_init=True)
   sp.add_argument("--gc",
                   action="store_true",
                   help="run garbage collection first")
   sp.add_argument("--reset",
                   action="store_true",
                   help="clear and re-initialize first")
   sp.add_argument("--tree",
                   action="store_true",
                   help="print a text tree summary")
   sp.add_argument("--dot", nargs="?", metavar="PATH", const="build-cache",
                   help="write DOT and PDF tree summaries")

   # delete
   sp = ap.add_parser("delete", "delete image from internal storage")
   add_opts(sp, misc.delete, deps_check=True, stog_init=True)
   sp.add_argument("image_ref", metavar="IMAGE_GLOB", help="image(s) to delete", nargs='+')

   # gestalt (has sub-subcommands)
   sp = ap.add_parser("gestalt", "query debugging configuration",
                      sub_title="subsubcommands", sub_metavar="CMD")
   add_opts(sp, lambda x: False, deps_check=False, stog_init=False)
   # bucache
   tp = sp.add_parser("bucache", "exit successfully if build cache available")
   add_opts(tp, misc.gestalt_bucache, deps_check=True, stog_init=False)
   # bucache-dot
   tp = sp.add_parser("bucache-dot", "exit success if can produce DOT trees")
   add_opts(tp, misc.gestalt_bucache_dot, deps_check=True, stog_init=False)
   # storage-path
   tp = sp.add_parser("storage-path", "print storage directory path")
   add_opts(tp, misc.gestalt_storage_path, deps_check=False, stog_init=False)
   # python-path
   tp = sp.add_parser("python-path", "print path to python interpreter in use")
   add_opts(tp, misc.gestalt_python_path, deps_check=False, stog_init=False)
   # logging
   tp = sp.add_parser("logging", "print logging messages at all levels")
   add_opts(tp, misc.gestalt_logging, deps_check=False, stog_init=False)
   tp.add_argument("--fail", action="store_true",
                   help="also generate a fatal error")

   # import
   sp = ap.add_parser("import", "copy external image into storage")
   add_opts(sp, misc.import_, deps_check=True, stog_init=True)
   sp.add_argument("path", metavar="PATH",
                   help="directory or tarball to import")
   sp.add_argument("image_ref", metavar="IMAGE_REF",
                   help="destination image name (tag)")

   # list
   sp = ap.add_parser("list", "print information about image(s)")
   add_opts(sp, misc.list_, deps_check=True, stog_init=True)
   sp.add_argument("-l", "--long", action="store_true",
                   help="use long listing format")
   sp.add_argument("-u", "--undeletable", action="store_true",
                   help="list images that can be restored with “undelete”")
   sp.add_argument("--undeleteable", action="store_true", dest="undeletable",
                   help=argparse.SUPPRESS)
   sp.add_argument("image_ref", metavar="IMAGE_REF", nargs="?",
                   help="print details of this image only")

   # modify
   sp = ap.add_parser("modify", "foo")
   add_opts(sp, modify.main, deps_check=True, stog_init=True)
   sp.add_argument("-c", metavar="CMD", action="append", default=[], nargs=1,
                   help="Run CMD as though specified by a RUN instruction. Can be repeated.")
   sp.add_argument("-i", "--interactive", action="store_true",
                   help="modify in interactive mode, even if stdin is not a TTY")
   sp.add_argument("-S", "--shell", metavar="shell", default="/bin/sh",
                   help="use SHELL instead of the default /bin/sh")
   sp.add_argument("image_ref", metavar="IMAGE_REF", help="image to modify")
   sp.add_argument("out_image", metavar="OUT_IMAGE", help="destination of modified image")
   sp.add_argument("script", metavar="SCRIPT", help="foo", nargs='?')
   # Options “modify” shares with “build”. Note that while we could abstract
   # this out to avoid repeated lines, as we do for “common_opts”, we’ve decided
   # that the tradeoff in code readability wouldn’t be worth it.
   sp.add_argument("-b", "--bind", metavar="SRC[:DST]",
                   action="append", default=[],
                   help="mount SRC at guest DST (default: same as SRC)")
   sp.add_argument("--build-arg", metavar="ARG[=VAL]",
                   action="append", default=[],
                   help="set build-time variable ARG to VAL, or $ARG if no VAL")
   sp.add_argument("--force", metavar="MODE", nargs="?", default="seccomp",
                   type=ch.Force_Mode, const="seccomp",
                   help="inject unprivileged build workarounds")
   sp.add_argument("--force-cmd", metavar="CMD,ARG1[,ARG2...]",
                   action="append", default=[],
                   help="command arg(s) to add under --force=seccomp")

   # pull
   sp = ap.add_parser("pull",
                      "copy image from remote repository to local filesystem")
   add_opts(sp, pull.main, deps_check=True, stog_init=True)
   sp.add_argument("--last-layer", metavar="N", type=int,
                   help="stop after unpacking N layers")
   sp.add_argument("--parse-only", action="store_true",
                   help="stop after parsing the image reference(s)")
   sp.add_argument("source_ref", metavar="IMAGE_REF", help="image reference")
   sp.add_argument("dest_ref", metavar="DEST_REF", nargs="?",
                   help="destination image reference (default: IMAGE_REF)")

   # push
   sp = ap.add_parser("push",
                      "copy image from local filesystem to remote repository")
   add_opts(sp, push.main, deps_check=True, stog_init=True)
   sp.add_argument("--image", metavar="DIR", type=fs.Path,
                   help="path to unpacked image (default: opaque path in storage dir)")
   sp.add_argument("source_ref", metavar="IMAGE_REF", help="image to push")
   sp.add_argument("dest_ref", metavar="DEST_REF", nargs="?",
                   help="destination image reference (default: IMAGE_REF)")

   # reset
   sp = ap.add_parser("reset", "delete everything in ch-image builder storage")
   add_opts(sp, misc.reset, deps_check=True, stog_init=False)

   # undelete
   sp = ap.add_parser("undelete", "recover image from build cache")
   add_opts(sp, misc.undelete, deps_check=True, stog_init=True)
   sp.add_argument("image_ref", metavar="IMAGE_REF", help="image to recover")

   # Monkey patch problematic characters out of stdout and stderr.
   ch.monkey_write_streams()

   # Parse it up!
   if (len(sys.argv) < 2):
       ap.print_help(file=sys.stderr)
       ch.exit(1)
   cli = ap.parse_args()

   # Initialize.
   ch.init(cli)
   if (dependencies_check[cli.func]):
      ch.dependencies_check()
   if (storage_init[cli.func]):
      ch.storage.init()
      bu.init(cli)

   # Dispatch.
   ch.profile_start()
   cli.func(cli)
   ch.warnings_dump()
   ch.exit(0)


## Functions ##

def breakpoint_inject(module_name, line_no):
   # Inject a PDB breakpoint into the module named module_name before the
   # statement on line line_no. See: https://stackoverflow.com/a/41858422

   class PDB_Injector(ast.NodeTransformer):
      def __init__(self, *args, **kwargs):
         self.inject_ct = 0
         return super().__init__(*args, **kwargs)
      def generic_visit(self, parent):
         # Operate on parent of target statement because we need to inject the
         # new code into the parent’s body (i.e., as siblings of the target
         # statement).
         if (    self.inject_ct == 0
             and hasattr(parent, "body")
             and isinstance(parent.body, collections.abc.Sequence)):
            for (i, child) in enumerate(parent.body):
               if (    isinstance(child, ast.stmt)
                   and hasattr(child, "lineno")
                   and child.lineno == line_no):
                  ch.WARNING(  "--break: injecting PDB breakpoint: %s:%d (%s)"
                             % (module_name, line_no, type(child).__name__))
                  parent.body[i:i] = inject_tree.body
                  self.inject_ct += 1
                  break
         super().generic_visit(parent)  # superclass actually visits children
         return parent

   if (module_name not in sys.modules):
      ch.FATAL("--break: no module named %s" % module_name)
   module = sys.modules[module_name]
   src_text = inspect.getsource(module)
   src_path = inspect.getsourcefile(module)
   module_tree = ast.parse(src_text, "%s <re-parsed>" % src_path)
   inject_tree = ast.parse("import pdb; pdb.set_trace()", "Weird Al Yankovic")

   ijor = PDB_Injector()
   ijor.visit(module_tree)  # calls generic_visit() on all nodes
   if (ijor.inject_ct < 1):
      ch.FATAL("--break: no statement found at %s:%d" % (module_name, line_no))
   assert (ijor.inject_ct == 1)

   ast.fix_missing_locations(module_tree)
   exec(compile(module_tree, "%s <re-compiled>" % src_path, "exec"),
        module.__dict__)
   # Set a global in the target module so it can test if it’s been
   # re-executed. This means re-execution is *complete*, so it will not be set
   # in module-level code run during re-execution, but if the original
   # execution continues *after* re-execution completes (this happens for
   # __main__), it *will* be set in that code.
   module.__dict__["breakpoint_reexecuted"] = "%s:%d" % (module_name, line_no)


## Bootstrap ##

# This code is more complicated than the standard boilerplace (i.e., “if
# (__name__ == "__main__"): main()”) for two reasons:
#
#   1. The mechanism for fatal errors is to raise ch.Fatal_Error. We catch
#      this to re-print warnings and print the error message before exiting.
#      (We used to priont an error message and then sys.exit(1), but this
#      approach lets us do things like rollback and fixes ordering problems
#      such as #1486.)
#
#   2. There is a big mess of hairy code to let us set PDB breakpoints in this
#      file (i.e., module __main__) with --break. See PR #1837.

if (__name__ == "__main__"):
   try:
      # We can’t set these two module globals that support --break normally
      # (i.e., module-level code at the top of this file) because this module
      # might be executed twice, and thus any value we set would be
      # overwritten by the default when the module is re-executed.
      if ("breakpoint_considered" not in globals()):
         global breakpoint_considered
         breakpoint_considered = True
         # A few lines of bespoke CLI parsing so that we can inject
         # breakpoints into the CLI parsing code itself.
         for (opt, arg) in zip(sys.argv[1:], sys.argv[2:] + [None]):
            (opt, _, arg_eq) = opt.partition("=")
            if (opt == "--break"):
               if (not sys.stdin.isatty()):
                  ch.FATAL("--break: standard input must be a terminal")
               if (arg_eq != ""):
                  arg = arg_eq
               try:
                  (module_name, line_no) = arg.split(":")
                  line_no = int(line_no)
               except ValueError:
                  ch.FATAL("--break: can’t parse MODULE:LIST: %s" % arg)
               breakpoint_inject(module_name, line_no)
      # If we injected into __main__, we already ran main() when re-executing
      # this module inside breakpoint_inject().
      if ("breakpoint_reexecuted" not in globals()):
         main()
   except ch.Fatal_Error as x:
      ch.warnings_dump()
      ch.ERROR(*x.args, **x.kwargs)
      ch.exit(1)
