#!/usr/bin/env python3

# Copyright (C) 2020  Braiins Systems s.r.o.
#
# This file is part of Braiins Open-Source Initiative (BOSI).
#
# BOSI is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
# Please, keep in mind that we may also license BOSI or any part thereof
# under a proprietary license. For more information on the terms and conditions
# of such proprietary license or if you have any other questions, please
# contact us at opensource@braiins.com.

import argparse
import concurrent.futures
import csv
import logging
import sys
import threading
import warnings

# these are needed for ssh method
from bos_toolbox.batch import read_hosts
from bos_utils.ssh import SSHManager
from subprocess import CalledProcessError

from .config import BosApi, UnsupportedVersion

LOG = logging.getLogger(__name__)


def main(parser, args):
    hosts = read_hosts(args.hosts)

    if args.command == ['start']:
        command = '/etc/init.d/bosminer start'
    elif args.command == ['stop']:
        command = '/etc/init.d/bosminer stop'
    elif args.command == ['restart']:
        command = '/etc/init.d/bosminer restart'
    else:
        command = ' '.join(args.command)

    runner = catcher_logger(args.runner)

    if args.jobs <= 1:
        # avoid executor swalowing exceptions
        for host in hosts:
            runner(host, 'root', args.password, command, args)
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.jobs) as executor:
            for host in hosts:
                executor.submit(runner, host, 'root', args.password, command, args)


def catcher_logger(fce):
    def catcher_logger_wrapper(*args, **kargs):
        try:
            fce(*args, **kargs)
        except Exception as ex:
            LOG.error(str(ex))

    return catcher_logger_wrapper


printer_lock = threading.Lock()


def run_command_auto(host, user, password, command, args):
    try:
        run_command_rpc(host, user, password, command, args)
    except UnsupportedVersion:
        run_command_ssh(host, user, password, command, args)


# @catcher_logger
def run_command_rpc(host, user, password, command, args):
    api = BosApi(host, user, password)
    # NOTE: this call does not return stderr
    stdout = api.rpc('sys', 'exec', command)

    with printer_lock:
        if args.print_host:
            LOG.info(host)
        if args.output:
            print(stdout, end='')


# @catcher_logger
def run_command_ssh(host, user, password, command, args):
    with SSHManager(host, user, password, load_host_keys=False) as client:
        try:
            stdout, stderr = client.run(command)
        except CalledProcessError as ex:
            stdout = ex.stdout
            stderr = ex.stderr
        with printer_lock:
            if args.print_host:
                LOG.info(host)
            if args.output:
                print(
                    stdout.read().decode('latin1'), end='', file=sys.stdout, flush=True
                )
                print(
                    stderr.read().decode('latin1'), end='', file=sys.stderr, flush=True
                )


def build_arg_parser(parser):
    parser.description = (
        'Run commands on one miner with hostname or IP or all miners in csv file. These may be start and stop '
        'for miner control or arbitrary shell commands'
    )

    parser.add_argument(
        'hosts',
        help='hostname or IP or path to file with hosts of miners with Braiins OS',
    )
    parser.add_argument(
        'command', nargs='+', help="command to run ('start', 'stop' or shell command)"
    )
    runner = parser.add_mutually_exclusive_group()
    runner.add_argument(
        '-a',
        '--auto',
        dest='runner',
        action='store_const',
        const=run_command_auto,
        default=run_command_auto,
        help='use ssh if rpc is not available',
    )
    runner.add_argument(
        '-l',
        '--legacy',
        dest='runner',
        action='store_const',
        const=run_command_ssh,
        help='use ssh',
    )
    runner.add_argument(
        '-L',
        '--no-legacy',
        dest='runner',
        action='store_const',
        const=run_command_rpc,
        help='use rpc',
    )
    parser.add_argument(
        '-o', '--output', action='store_true', help='capture and print remote output'
    )
    parser.add_argument(
        '-O',
        '--output-hostname',
        dest='print_host',
        action='store_true',
        help='capture and print remote output',
    )
    parser.add_argument('-p', '--password', default='', help='administration password')
    parser.add_argument(
        '-j', '--jobs', type=int, default=50, help='number of concurrent jobs'
    )


if __name__ == '__main__':
    warnings.filterwarnings('ignore', module='paramiko')
    parser = argparse.ArgumentParser()
    build_arg_parser(parser)
    args = parser.parse_args()
    main(parser, args)
