"""Allows the package to be run with `python3 -m ubuntu_image`."""

import os
import sys
import logging
import argparse

from contextlib import suppress
from pickle import dump, load
from ubuntu_image import __version__
from ubuntu_image.builder import DoesNotFit, ModelAssertionBuilder
from ubuntu_image.helpers import as_size
from ubuntu_image.i18n import _
from ubuntu_image.parser import GadgetSpecificationError


_logger = logging.getLogger('ubuntu-image')


PROGRAM = 'ubuntu-image'


class SizeAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        sizes = {}
        specs = values.split(',')
        # First check for extended syntax by splitting on commas.
        for spec in specs:
            index, colon, size = spec.partition(':')
            if colon != ':':
                if len(specs) != 1:
                    raise argparse.ArgumentError(
                        self,
                        'Invalid multi-volume size specification: {}'.format(
                            spec))
                # Backward compatibility.  Since there was no colon to
                # partition the string on, the size argument will be in the
                # index local variable.
                try:
                    sizes = as_size(index)
                except (KeyError, ValueError):
                    raise argparse.ArgumentError(
                        self, 'Invalid size: {}'.format(values))
                break
            try:
                size = as_size(size)
            except (KeyError, ValueError):
                raise argparse.ArgumentError(
                    self, 'Invalid size: {}'.format(values))
            try:
                index = int(index)
            except (ValueError, TypeError):
                pass
            sizes[index] = size
        setattr(namespace, self.dest, sizes)
        # For display purposes.
        namespace.given_image_size = values


def parseargs(argv=None):
    parser = argparse.ArgumentParser(
        prog=PROGRAM,
        description=_('Generate a bootable disk image.'),
        )
    parser.add_argument(
        '--version', action='version',
        version='{} {}'.format(PROGRAM, __version__))
    # Common options.
    common_group = parser.add_argument_group(_('Common options'))
    common_group.add_argument(
        'model_assertion', nargs='?',
        help=_("""Path to the model assertion file.  This argument must be
        given unless the state machine is being resumed, in which case it
        cannot be given."""))
    common_group.add_argument(
        '-d', '--debug',
        default=False, action='store_true',
        help=_('Enable debugging output'))
    common_group.add_argument(
        '-i', '--image-size',
        default=None, action=SizeAction, metavar='SIZE',
        help=_("""The suggested size of the generated disk image file.  If this
        size is smaller than the minimum calculated size of the image a warning
        will be issued and --image-size will be ignored.  The value is the size
        in bytes, with allowable suffixes 'M' for MiB and 'G' for GiB.  Use an
        extended syntax to define the suggested size for the disk images
        generated by a multi-volume gadget.yaml spec.  See the ubuntu-image(1)
        manpage for details."""))
    common_group.add_argument(
        '--image-file-list',
        default=None, metavar='FILENAME',
        help=_("""Print to this file, a list of the file system paths to
        all the disk images created by the command, if any."""))
    output_group = common_group.add_mutually_exclusive_group()
    output_group.add_argument(
        '-O', '--output-dir',
        default=None, metavar='DIRECTORY',
        help=_("""The directory in which to put generated disk image files.
        The disk image files themselves will be named <volume>.img inside this
        directory, where <volume> is the volume name taken from the
        gadget.yaml file.  Use this option instead of the deprecated
        -o/--output option."""))
    output_group.add_argument(
        '-o', '--output',
        default=None, metavar='FILENAME',
        help=_("""DEPRECATED (use -O/--output-dir instead).  The generated
        disk image file.  If not given, the image will be put in a file called
        disk.img in the working directory (in which case, you probably want to
        specify -w)."""))
    # Snap-based image options.
    snap_group = parser.add_argument_group(
        _('Image contents options'),
        _("""Additional options for defining the contents of snap-based
        images."""))
    snap_group.add_argument(
        '--extra-snaps',
        default=None, action='append',
        help=_("""Extra snaps to install.  This is passed through to `snap
        prepare-image`."""))
    snap_group.add_argument(
        '--cloud-init',
        default=None, metavar='USER-DATA-FILE',
        help=_('cloud-config data to be copied to the image'))
    snap_group.add_argument(
        '-c', '--channel',
        default=None,
        help=_('The snap channel to use'))
    # State machine options.
    inclusive_state_group = parser.add_argument_group(
        _('State machine options'),
        _("""Options for controlling the internal state machine.  Other than
        -w, these options are mutually exclusive.  When -u or -t is given, the
        state machine can be resumed later with -r, but -w must be given in
        that case since the state is saved in a .ubuntu-image.pck file in the
        working directory."""))
    inclusive_state_group.add_argument(
        '-w', '--workdir',
        default=None, metavar='DIRECTORY',
        help=_("""The working directory in which to download and unpack all
        the source files for the image.  This directory can exist or not, and
        it is not removed after this program exits.  If not given, a temporary
        working directory is used instead, which *is* deleted after this
        program exits.  Use -w if you want to be able to resume a partial
        state machine run."""))
    state_group = inclusive_state_group.add_mutually_exclusive_group()
    state_group.add_argument(
        '-u', '--until',
        default=None, metavar='STEP',
        help=_("""Run the state machine until the given STEP, non-inclusively.
        STEP can be a name or number."""))
    state_group.add_argument(
        '-t', '--thru',
        default=None, metavar='STEP',
        help=_("""Run the state machine through the given STEP, inclusively.
        STEP can be a name or number."""))
    state_group.add_argument(
        '-r', '--resume',
        default=False, action='store_true',
        help=_("""Continue the state machine from the previously saved state.
        It is an error if there is no previous state."""))
    args = parser.parse_args(argv)
    if args.debug:
        logging.basicConfig(level=logging.DEBUG)
    # The model assertion argument is required unless --resume is given, in
    # which case it cannot be given.
    if args.resume and args.model_assertion:
        parser.error('model assertion is not allowed with --resume')
    if not args.resume and args.model_assertion is None:
        parser.error('model assertion is required')
    if args.resume and args.workdir is None:
        parser.error('--resume requires --workdir')
    # --until and --thru can take an int.
    with suppress(ValueError, TypeError):
        args.thru = int(args.thru)
    with suppress(ValueError, TypeError):
        args.until = int(args.until)
    # -o/--output is deprecated and mutually exclusive with -O/--output-dir
    if args.output is not None:
        print('-o/--output is deprecated; use -O/--output-dir instead',
              file=sys.stderr)
    return args


def main(argv=None):
    args = parseargs(argv)
    if args.workdir:
        os.makedirs(args.workdir, exist_ok=True)
        pickle_file = os.path.join(args.workdir, '.ubuntu-image.pck')
    else:
        pickle_file = None
    if args.resume:
        with open(pickle_file, 'rb') as fp:
            state_machine = load(fp)
        state_machine.workdir = args.workdir
    else:
        state_machine = ModelAssertionBuilder(args)
    # Run the state machine, either to the end or thru/until the named state.
    try:
        if args.thru:
            state_machine.run_thru(args.thru)
        elif args.until:
            state_machine.run_until(args.until)
        else:
            list(state_machine)
    except GadgetSpecificationError as error:
        if args.debug:
            _logger.exception('gadget.yaml parse error')
        else:
            _logger.error('gadget.yaml parse error: {}'.format(error))
            _logger.error('Use --debug for more information')
    except DoesNotFit as error:
        _logger.error(
            'Volume contents do not fit ({}B over): {} [#{}]'.format(
                error.overage, error.part_path, error.part_number))
    except:
        _logger.exception('Crash in state machine')
        return 1
    # It's possible that the state machine didn't crash, but it still didn't
    # complete successfully.  For example, if `snap prepare-image` failed.
    if state_machine.exitcode != 0:
        return state_machine.exitcode
    # Write out the list of images, if there are any.
    if (state_machine.gadget is not None and
            state_machine.done and
            args.image_file_list is not None):
        with open(args.image_file_list, 'w', encoding='utf-8') as fp:
            if args.output is None:
                for name in state_machine.gadget.volumes:
                    path = os.path.join(
                        args.output_dir, '{}.img'.format(name))
                    print(path, file=fp)
            else:
                print(args.output, file=fp)
    # Everything's done, now handle saving state if necessary.
    if pickle_file is not None:
        with open(pickle_file, 'wb') as fp:
            dump(state_machine, fp)
    return 0


if __name__ == '__main__':                          # pragma: nocover
    sys.exit(main())
