#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (C) 2018-2022 fjord-technologies
# SPDX-License-Identifier: GPL-3.0-or-later
"""auton"""

__version__ = '1.0.0'

import argparse
import os
import shlex
import sys
import uuid

import logging
import math
import json
from logging.handlers import WatchedFileHandler

from dotenv.main import dotenv_values
from sonicprobe import helpers

import requests  # Compatibility for clients patching the historical transport.
from auton_client import RemoteClient
from auton_client.credentials import credentials_from_options

import six

SYSLOG_NAME     = "auton"
LOG             = logging.getLogger(SYSLOG_NAME)

DEFAULT_LOGFILE = "/var/log/auton/auton.log"
DEFAULT_DELAY   = 0.5

AUTON_LOGFILE   = os.environ.get('AUTON_LOGFILE') or DEFAULT_LOGFILE


def argv_parse_check():
    """
    Parse (and check a little) command line parameters
    """
    parser        = argparse.ArgumentParser()

    parser.add_argument('-c', '--config', help='Client YAML inventory; targets must still be selected explicitly')
    parser.add_argument('--tui', action='store_true', help='Read-only ncurses operator interface')
    parser.add_argument('--daemon', action='append', default=[], metavar='SELECTOR[=URI]',
                        help='TUI target name/glob/~regex from --config or NAME=URI')
    parser.add_argument('-t', '--target', action='append', default=[], metavar='SELECTOR[=URI]',
                        help='Target name/glob/~regex from --config or NAME=URI')
    parser.add_argument('-g', '--target-group', action='append', default=[], metavar='SELECTOR',
                        help='Select YAML groups by name, glob or ~regex; repeat to combine')
    parser.add_argument('-s', '--scenario', action='append', default=[], metavar='SELECTOR',
                        help='Execute named YAML scenarios, selected by name/glob/~regex')
    parser.add_argument('-S', '--scenario-group', action='append', default=[], metavar='SELECTOR',
                        help='Execute ordered YAML scenario groups, selected by name/glob/~regex')
    parser.add_argument('--parallel', type=int, default=4, help='Maximum concurrent targets (1-32)')
    parser.add_argument('--operation-timeout', type=float, default=300,
                        help='Stop observing after this many seconds; does not cancel remote jobs')
    parser.add_argument('--operation-id', help='Optional operation label (client-side only)')
    parser.add_argument('--refresh', type=float, default=2.0,
                        help='TUI refresh interval in seconds (minimum 0.2, default: 2)')

    parser.add_argument("-a",
                        action    = 'append',
                        dest      = 'args',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Passed arguments")
    parser.add_argument("-A",
                        action    = 'append',
                        dest      = 'argfiles',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Passed argument files")
    parser.add_argument("--multi-args",
                        action    = 'append',
                        dest      = 'margs',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Passed multiple arguments (arguments splitted)")
    parser.add_argument("--multi-argsfiles",
                        action    = 'append',
                        dest      = 'margsfiles',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Passed multiple arguments files (arguments splitted)")
    parser.add_argument("--uri",
                        action    = 'append',
                        dest      = 'uri',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Auton daemon URI addresses")
    parser.add_argument("--uid",
                        dest      = 'uid',
                        type      = six.ensure_text,
                        default   = os.environ.get('AUTON_UID') or None,
                        help      = "Auton uid")
    parser.add_argument("--endpoint",
                        dest      = 'endpoint',
                        type      = six.ensure_text,
                        default   = os.environ.get('AUTON_ENDPOINT'),
                        help      = "Auton endpoint")
    parser.add_argument('-k', '--token-file', default=os.environ.get('AUTON_TOKEN_FILE'),
                        help='Private bearer-token file; mutually exclusive with Basic credentials')
    parser.add_argument("--auth-user",
                        dest      = 'auth_user',
                        type      = six.ensure_text,
                        default   = os.environ.get('AUTON_AUTH_USER'),
                        help      = "Auton auth user")
    parser.add_argument("--auth-passwd",
                        dest      = 'auth_passwd',
                        type      = six.ensure_text,
                        default   = os.environ.get('AUTON_AUTH_PASSWD'),
                        help      = "Auton auth password")
    parser.add_argument("--delay",
                        dest      = 'delay',
                        type      = float,
                        default   = DEFAULT_DELAY,
                        help      = "Delay between requests instead of %(default)s")
    parser.add_argument('--http-timeout', type=float, default=30,
                        help='HTTP connect/read timeout in seconds (default: 30)')
    parser.add_argument("--mode",
                        dest      = 'mode',
                        default   = 'autorun',
                        choices   = ('autorun', 'run', 'status'),
                        help      = "Auton mode: autorun, run, status, instead of %(default)s")
    parser.add_argument("-e",
                        action    = 'append',
                        dest      = 'envvars',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Passed environment variables")
    parser.add_argument("--envfile",
                        action    = 'append',
                        dest      = 'envfiles',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Passed envfile parameters")
    parser.add_argument("--imp-env",
                        action    = 'append',
                        dest      = 'imp_envvars',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Import existing environment variables")
    parser.add_argument("--load-envfile",
                        action    = 'append',
                        dest      = 'load_envfiles',
                        type      = six.ensure_text,
                        default   = [],
                        help      = "Load environment variables from file")
    parser.add_argument("-l",
                        dest      = 'loglevel',
                        default   = 'info',   # warning: see affectation under
                        choices   = ('critical', 'error', 'warning', 'info', 'debug'),
                        help      = ("Emit traces with LOGLEVEL details, must be one of:\t"
                                     "critical, error, warning, info, debug"))
    parser.add_argument("--logfile",
                        dest      = 'logfile',
                        type      = six.ensure_text,
                        default   = AUTON_LOGFILE,
                        help      = "Use log file <logfile> instead of %(default)s")
    parser.add_argument("--no-return-code",
                        action    = 'store_true',
                        dest      = 'no_return_code',
                        default   = helpers.boolize(os.environ.get('AUTON_NO_RETURN_CODE', False)),
                        help      = "Do not exit with return code if present")

    sys.argv      = helpers.escape_parse_args(('-a',
                                               '-A',
                                               '--multi-args',
                                               '--multi-argsfiles'),
                                              sys.argv)

    options, args = parser.parse_known_args() # pylint: disable=unused-variable
    options.loglevel = getattr(logging, options.loglevel.upper(), logging.INFO)
    if (not math.isfinite(options.http_timeout) or options.http_timeout <= 0
            or not math.isfinite(options.delay) or options.delay < 0):
        parser.error('http-timeout must be positive and delay non-negative')

    if not options.uri and not options.daemon and not options.target and not options.target_group and not (options.tui and options.config) and os.environ.get('AUTON_URI'):
        options.uri = [x.strip() for x in os.environ.get('AUTON_URI').split(',')]

    if options.target or options.target_group:
        if options.uri or (options.daemon and options.target) or (not options.tui and
                (options.mode != 'autorun' or options.uid is not None)):
            parser.error('target selections cannot use --uri, --uid/AUTON_UID, --mode run/status or mix --target with --daemon')
        if options.no_return_code and not options.tui:
            parser.error('--no-return-code cannot hide multi-target failures')
    if (options.tui or not (options.target or options.target_group)) and (options.operation_id or options.parallel != 4 or options.operation_timeout != 300):
        parser.error('operation options require execution with --target or --target-group')
    if options.uid is None:
        options.uid = str(uuid.uuid4())

    if options.daemon and not options.tui:
        parser.error('--daemon requires --tui; it is not a broadcast execution option')
    if options.tui:
        if not math.isfinite(options.refresh) or options.refresh < 0.2:
            parser.error('refresh must be at least 0.2 seconds')
        if (options.mode != 'autorun' or options.args or options.argfiles or options.margs
                or options.margsfiles or options.envvars or options.envfiles
                or options.imp_envvars or options.load_envfiles):
            parser.error('--tui prepares execution interactively; CLI payload options are incompatible')

    scenario_mode = bool(options.scenario or options.scenario_group)
    if scenario_mode:
        if not options.config or not (options.tui or options.target or options.target_group):
            parser.error('scenario execution requires --config and explicit target selections')
        if (options.endpoint or options.mode != 'autorun' or options.args
                or options.argfiles or options.margs or options.margsfiles or options.envvars
                or options.envfiles or options.imp_envvars or options.load_envfiles):
            parser.error('scenario definitions own endpoints and inputs; incompatible execution options')

    if options.token_file and (options.auth_user or options.auth_passwd):
        parser.error('token-file and Basic credentials are mutually exclusive')

    options.selected_scenarios = None
    options.configured_targets = {}
    options.configured_credentials = {}
    options.configured_groups = {}
    options.configured_scenarios = {}
    options.configured_scenario_groups = {}
    if options.target_group and not options.config:
        parser.error('--target-group requires --config')
    if options.config:
        if options.uri or not (options.tui or options.target or options.daemon or options.target_group):
            parser.error('--config requires explicit target/group/daemon selections and cannot use --uri')
        try:
            from auton_client.config import load_inventory
            inventory = load_inventory(options.config)
            options.configured_targets = inventory['targets']
            options.configured_credentials = inventory['credentials']
            options.configured_groups = inventory['groups']
            options.configured_scenarios = inventory['scenarios']
            options.configured_scenario_groups = inventory['scenario_groups']
            if scenario_mode:
                from auton_client.scenarios import select_scenarios
                options.selected_scenarios = select_scenarios(options.scenario, options.scenario_group,
                    inventory['scenarios'], inventory['scenario_groups'])
        except ValueError as error:
            parser.error(str(error))
    options.selected_targets = None
    if options.tui and options.config:
        options.selected_targets = options.configured_targets
    if options.target or options.daemon or options.target_group:
        try:
            from auton_client.connections import select_connections
            options.selected_targets = select_connections(options.target or options.daemon,
                options.target_group, options.configured_targets, options.configured_groups)
            if not options.selected_targets:
                raise ValueError("selection contains no targets")
        except ValueError as error:
            parser.error(str(error))

    return options


class PayloadBuilder:
    def __init__(self, options):
        self.options = options
        self.envvars = {}
        self.argfiles = []
        self._parse_multi_args()
        self._parse_multi_argsfiles()
        self._import_envvars()
        self._load_envfiles()
        self._parse_envvars()
        self._parse_argfiles()
        self.payload = {'args': options.args, 'argfiles': self.argfiles,
                        'env': self.envvars, 'envfiles': options.envfiles}

    def _import_envvars(self):
        for envvar in self.options.imp_envvars:
            if envvar in os.environ:
                self.envvars[envvar] = os.environ[envvar]
            else:
                LOG.debug("unable to find environment variable: %r", envvar)

    def _load_envfiles(self):
        for envfile in self.options.load_envfiles:
            self.envvars.update(dotenv_values(envfile))

    def _parse_envvars(self):
        for envvar in self.options.envvars:
            env = envvar.split('=', 1)
            if not env[0]:
                LOG.warning("invalid environment variable: %r", env)
            elif len(env) == 1:
                self.envvars[env[0]] = ''
            else:
                self.envvars[env[0]] = env[1]

    def _parse_multi_args(self):
        for margs in self.options.margs:
            args = shlex.split(margs)
            if args:
                self.options.args.extend(args)

    def _parse_multi_argsfiles(self):
        for margsfiles in self.options.margsfiles:
            argsfiles = shlex.split(margsfiles)
            if argsfiles:
                self.options.argfiles.extend(argsfiles)

    def _parse_argfiles(self):
        for argfile in self.options.argfiles:
            arg = argfile.split('=', 1)
            if not arg[0] or arg[0] == '@' or len(arg) == 1:
                LOG.warning("invalid argument file name: %r", arg)
            elif arg[1] == '-':
                data = six.BytesIO(helpers.read_large_file(sys.stdin))
                self.argfiles.append({'arg': arg[0],
                                      'content': helpers.base64_encode_file(data),
                                      'filename': ''})
                data = None
            elif not os.path.isfile(arg[1]):
                LOG.warning("unable to find file: %r", arg[1])
            else:
                self.argfiles.append({'arg': arg[0],
                                      'content': helpers.base64_encode_file(arg[1]),
                                      'filename': os.path.basename(arg[1])})


class AutonClient(RemoteClient, PayloadBuilder):
    def __init__(self, options):
        if not options.uri:
            raise ValueError('missing variable AUTON_URI')
        if not options.uid:
            raise ValueError('missing variable AUTON_UID')
        if not options.endpoint:
            raise ValueError('missing variable AUTON_ENDPOINT')
        PayloadBuilder.__init__(self, options)
        auth = credentials_from_options(options)
        RemoteClient.__init__(self, options.uri, options.endpoint, options.uid,
                              payload=self.payload, auth=auth, http_timeout=options.http_timeout)

    @staticmethod
    def _show_results(data):
        for chunk in data.get('stream') or []:
            sys.stdout.write(chunk)
        if data['status'] == 'complete':
            for chunk in data.get('errors') or []:
                sys.stderr.write(chunk)

    def do_autorun(self):
        for data in self.iter_results(delay=self.options.delay):
            self._show_results(data)
        if self.options.no_return_code or data['return_code'] is None:
            return int(data['code'] != 200)
        return data['return_code']


def main(options):
    """
    Main function
    """
    if options.tui:
        try:
            from auton_client.tui import run
            auth = credentials_from_options(options)
            return_code = run(options.daemon, options.uri, auth, options.http_timeout, options.refresh,
                              configured=options.configured_targets, selected=options.selected_targets,
                              groups=options.configured_groups,
                              scenarios=options.selected_scenarios if options.selected_scenarios is not None else options.configured_scenarios,
                              scenario_groups=options.configured_scenario_groups)
        except ImportError:
            print('Auton TUI requires Python curses support.', file=sys.stderr)
            return_code = 1
        except KeyboardInterrupt:
            return_code = 2
        except ValueError as error:
            print(str(error), file=sys.stderr)
            return_code = 1
        sys.exit(return_code)

    xformat = "%(levelname)s:%(asctime)-15s: %(message)s"
    datefmt = '%Y-%m-%d %H:%M:%S'
    logging.basicConfig(level   = options.loglevel,
                        format  = xformat,
                        datefmt = datefmt)

    logdir  = os.path.dirname(options.logfile)
    if os.path.isdir(logdir) and os.access(logdir, os.W_OK):
        filehandler = WatchedFileHandler(options.logfile)
        filehandler.setFormatter(logging.Formatter(xformat,
                                                   datefmt = datefmt))
        root_logger = logging.getLogger('')
        root_logger.addHandler(filehandler)

    auton   = None
    rc      = 0

    try:
        if options.target or options.target_group:
            from auton_client.operations import OperationService
            auth = credentials_from_options(options)
            settings = dict(auth=auth, http_timeout=options.http_timeout, parallel=options.parallel,
                            timeout=options.operation_timeout, delay=options.delay)
            if options.selected_scenarios:
                from auton_client.scenarios import ScenarioService
                operation = ScenarioService(options.selected_targets, options.selected_scenarios, **settings)
            else:
                operation = OperationService(options.selected_targets, options.endpoint,
                                             payload=PayloadBuilder(options).payload, **settings)
            result = operation.run(options.operation_id)
            print(json.dumps(result))
            rc = int(result['status'] != 'completed')
        else:
            auton = AutonClient(options)
            rc = getattr(auton, "do_%s" % options.mode)()
        if isinstance(rc, dict):
            print(json.dumps(rc))
            rc = int(rc['code'] != 200)
    except (KeyboardInterrupt, SystemExit):
        rc = 2
    except Exception as e:
        LOG.error(e)
        rc = 1
        raise
    finally:
        sys.exit(rc)


if __name__ == '__main__':
    main(argv_parse_check())
