#!/usr/bin/python3

import argparse
import subprocess
import sys

# TODO: Better multi-ssh

def main():
  parser = argparse.ArgumentParser(description='Remote control.', add_help=False)
  parser.add_argument('-h', '--host-file', metavar='host-file', type=str, help="file with hosts, one per line")
  parser.add_argument('-H', '--host', metavar='host', type=str, action='append', help="additional hosts")
  parser.add_argument('-O', '--option', metavar='option', type=str, action='append', help="ssh option")
  subparsers = parser.add_subparsers()

  parser_become = subparsers.add_parser('become', help='become a service')
  parser_become.add_argument('--yaml', metavar='file', type=str, help='path to service.yaml')
  parser_become.add_argument('--type', metavar='type', type=str, help='service type, needed when using --yaml')
  parser_become.add_argument('--hash', metavar='hash', type=str, help='hash of previously deployed service')
  parser_become.set_defaults(func=do_become)

  parser_status = subparsers.add_parser('status', help='check service status')
  parser_status.add_argument('--output', type=str, choices=['short', 'long', 'json'], help='output style')
  parser_status.set_defaults(func=do_status)

  parser_status = subparsers.add_parser('history', help='read service history')
  parser_status.set_defaults(func=do_history)

  args = parser.parse_args()
  return args.func(args)

def do_become(args):
  exit_code = 0

  service_yaml = None
  if args.yaml:
    with open(args.yaml, "r") as f:
      service_yaml = f.read()

  for host in get_hosts(args):
    sys.stdout.write("[{}]\n".format(host))
    sys.stdout.flush()
    command = get_ssh_command(args) + [host, 'sudo', 'grove-become']
    if args.yaml:
      command = command + ['--yaml', '/dev/fd/0']
    if args.type:
      command = command + ['--type', args.type]
    if args.hash:
      command = command + ['--hash', args.hash]
    p = subprocess.Popen(command, stdin=subprocess.PIPE, stderr=subprocess.STDOUT)
    if args.yaml:
      p.stdin.write(service_yaml.encode())
    p.stdin.close()
    p.wait()

    if p.returncode != 0:
      sys.stdout.write("command failed: " + str(command) + "\n")
      sys.stdout.flush()
      exit_code = 1

  return exit_code

def do_status(args):
  exit_code = 0

  for host in get_hosts(args):
    sys.stdout.write("[{}]\n".format(host))
    sys.stdout.flush()
    command = get_ssh_command(args) + [host, 'sudo', 'grove-status']

    if args.output:
      command += ["--output", args.output]

    p = subprocess.Popen(command, stderr=subprocess.STDOUT)
    p.wait()

    if p.returncode != 0:
      sys.stdout.write("command failed: " + str(command) + "\n")
      sys.stdout.flush()
      exit_code = 1

  return exit_code

def do_history(args):
  exit_code = 0

  for host in get_hosts(args):
    sys.stdout.write("[{}]\n".format(host))
    sys.stdout.flush()
    command = get_ssh_command(args) + [host, 'sudo', 'grove-history']
    p = subprocess.Popen(command, stderr=subprocess.STDOUT)
    p.wait()

    if p.returncode != 0:
      sys.stdout.write("command failed: " + str(command) + "\n")
      sys.stdout.flush()
      exit_code = 1

  return exit_code

def get_ssh_command(args):
  command = ['ssh']
  if args.option is not None:
    for option in args.option:
      command = command + ['-o', option]
  return command

def get_hosts(args):
  hosts = []

  if args.host is not None:
    for host in args.host:
      hosts.append(host)

  if args.host_file is not None:
    with open(args.host_file, 'r') as f:
      for host in f:
        if not host.startswith('#') and not host.rstrip() == "":
          hosts.append(host.rstrip())

  return hosts

sys.exit(main())
