#!/usr/bin/env python

# Copyright 2014-present, Apstra, Inc. All rights reserved.
#
# This source code is licensed under End User License Agreement found in the
# LICENSE file at http://www.apstra.com/community/eula

import sys
import os
import json
import argparse
import yaml
import logging

import aeon.nxos as nxos
from aeon.nxos.exceptions import NxosException
from aeon.exceptions import ProbeError, UnauthorizedError

_OS_NAME = 'nxos'
_PROGNAME = '%s-installos' % _OS_NAME


# ##### -----------------------------------------------------------------------
# #####
# #####                           Command Line Arguments
# #####
# ##### -----------------------------------------------------------------------

psr = argparse.ArgumentParser(
    prog='nxos-installos',
    description="NXOS install OS",
    add_help=True)

psr.add_argument(
    '-t', '--target',
    required=True,
    help='hostname or ip_addr of target device')

psr.add_argument(
    '--server',
    required=True,
    help='ip-addr of the server hosting the image files')

psr.add_argument(
    '--image',
    required=True,
    help='image file-name')

psr.add_argument(
    '--md5sum', required=True,
    help='MD5 checksum value')

psr.add_argument(
    '--logfile',
    help='name of log file')

# ##### -------------------------
# ##### authentication
# ##### -------------------------

group = psr.add_argument_group('authentication')

group.add_argument(
    '--user', help='login user-name')

group.add_argument(
    '-U', dest='env_user',
    help='Username environment variable')

group.add_argument(
    '-P', dest='env_passwd',
    required=True,
    help='Passwd environment variable')

g_cli_args = psr.parse_args()


def setup_logging(logname, logfile, target):
    log = logging.getLogger(name=logname)
    log.setLevel(logging.INFO)

    fmt = logging.Formatter(
        '%(asctime)s:%(levelname)s:{target}:%(message)s'
        .format(target=target))

    handler = logging.FileHandler(logfile) if logfile else logging.StreamHandler(sys.stdout)
    handler.setFormatter(fmt)
    log.addHandler(handler)
    return log


g_log = setup_logging(logname=_PROGNAME,
                      logfile=g_cli_args.logfile,
                      target=g_cli_args.target)
def setup_logging():
    g_log.setLevel(logging.INFO)
    fh = logging.FileHandler(g_cli_args.logfile)
    fmt = logging.Formatter(
        '%(asctime)s:%(levelname)s: {target}:%(message)s'
        .format(target=g_cli_args.target))
    fh.setFormatter(fmt)
    g_log.addHandler(fh)

if g_cli_args.logfile:
    setup_logging()


# ##### -----------------------------------------------------------------------
# #####
# #####                           Utility Functions
# #####
# ##### -----------------------------------------------------------------------

def exit_results(results, exit_error=None):
    json.dump(results, fp=sys.stdout)
    sys.exit(0 if results['ok'] is True else exit_error or 1)


# ##### -----------------------------------------------------------------------
# #####
# #####                           Main Code Functions
# #####
# ##### -----------------------------------------------------------------------

def copy_image_to_device(dev):

    copy_from = "http://{server}/images/{os_name}/{image_file}".format(
        server=g_cli_args.server, os_name=_OS_NAME,
        image_file=g_cli_args.image)

    # --------------------------------------------------------
    # perform the copy, the device initiates this copy command
    # --------------------------------------------------------

    g_log.info('starting copy [%s] ... please be patient' % copy_from)
    try:
        dev.install_os.copy_from(location=copy_from)

    except NxosException:
        msg = 'failed to copy image: %s' % copy_from
        g_log.error(msg)
        exit_results(dict(
            ok=False,
            error_type='install',
            message=msg))

    # --------------------------------------------
    # now validate the MD5 checksum after the copy
    # --------------------------------------------

    has_md5sum = dev.install_os.md5sum
    if has_md5sum != g_cli_args.md5sum:
        msg = 'md5sum mismatch after file copy, need user attention'
        g_log.error(msg)
        exit_results(dict(
            ok=False,
            error_type='install',
            message=msg))
    else:
        g_log.info("md5sum check pass, copy completed OK.")


def start(dev):

    # TODO: add back the version string check before copy
    # req_os_version = profile['os_version']
    # dev_os_version = dev.facts['os_version']
    #
    # g_log.info("required OS version is: %s" % req_os_version)
    # g_log.info("device is current running: %s" % dev_os_version)
    #
    # if req_os_version == dev_os_version:
    #     msg = 'image is already installed, so all done.'
    #     g_log.info(msg)
    #     exit_results(dict(
    #         ok=True,
    #         rebooting=False,
    #         message=msg
    #     ))

    # -----------------------------------------------------------------
    # does the image exist on the device and/or need to be copied down?
    # -----------------------------------------------------------------

    dev.install_os(image=g_cli_args.image)
    has_md5sum = dev.install_os.md5sum

    if not has_md5sum:
        # need to copy the image over to the device
        g_log.info('image does not exist on device, need to copy over')
        copy_image_to_device(dev)

    elif has_md5sum != g_cli_args.md5sum:
        # need to make sure the image/md5sum matches
        # @@@ TO-DO: implement this code branch

        msg = 'existing file md5sum mismatch, need to re-copy image'
        g_log.warn(msg)
        exit_results(dict(
            ok=False,
            error_type='not-implemented',
            message=msg))
    else:
        g_log.info('image file exists and md5sum matches OK.')

    # -----------------------------------------------------------------
    # at this point, the image exists, and md5sum matches
    # now need to run the install process.  this will cause the device
    # to reboot when the installation is completed
    # -----------------------------------------------------------------

    g_log.info('running installation process ... please be patient')
    try:
        dev.install_os.run()
    except NxosException:
        msg = 'failed to install properly'
        g_log.error(msg)
        exit_results(dict(
            ok=False,
            error_type='install',
            message=msg))

    g_log.info('installation process completed, device rebooting')
    exit_results(dict(
        ok=True,
        rebooting=True,
        message='installation completed'))


def main():
    target = g_cli_args.target
    user = g_cli_args.user or os.getenv(g_cli_args.env_user)
    passwd = os.getenv(g_cli_args.env_passwd)

    if not user:
        errmsg = "login user-name missing"
        g_log.error(errmsg)
        exit_results(results=dict(
            ok=False,
            error_type='login',
            message=errmsg))

    if not passwd:
        errmsg = "login user-password missing"
        g_log.error(errmsg)
        exit_results(results=dict(
            ok=False,
            error_type='login',
            message=errmsg))

    try:
        dev = nxos.Device(target, user=user, passwd=passwd)

    except UnauthorizedError:
        exit_results(dict(
            ok=False,
            error_type='login',
            message='Uauthorized - check user/password'))

    except ProbeError:
        exit_results(dict(
            ok=False,
            error_type='login',
            message='Failed to probe target %s' % target))

    else:
        start(dev)

if '__main__' == __name__:
    main()
