#!/usr/bin/python3

import argparse
import fnmatch
import json
import os
import re
import socket
import sys
import time

def init():
  parser = argparse.ArgumentParser(description='Interact with the Grove server.')
  parser.add_argument('--url',
    type=str, default=os.environ.get('GROVE_SERVER_URL'), help='grove server url')
  parser.add_argument('--name', '-n', type=str, help='select service name glob')
  parser.add_argument('--type', '-t', type=str, help='select service type glob')
  parser.add_argument('--host', '-H', type=str, action='append', help='select host glob')
  parser.add_argument('--online', action='store_true', help='select online hosts')
  parser.add_argument('--offline', action='store_true', help='select offline hosts')
  parser.add_argument('--all', '-a', action='store_true', help='select all')
  subparsers = parser.add_subparsers(dest='command')
  subparsers.required = True

  parser_status = subparsers.add_parser('status', help='get agent status')
  parser_status.set_defaults(func='do_status')

  parser_show = subparsers.add_parser('show', help='get agent status')
  parser_show.set_defaults(func='do_status')

  parser_bounce = subparsers.add_parser('bounce', help='bounce services')
  parser_bounce.set_defaults(func='do_bounce')

  parser_reap = subparsers.add_parser('reap', help='reap offline agents')
  parser_reap.set_defaults(func='do_reap')

  parser_history = subparsers.add_parser('history', help='get agent deploy history')
  parser_history.set_defaults(func='do_history')

  parser_become = subparsers.add_parser('become', help='deploy software to agents')
  parser_become.add_argument('--yaml', metavar='file', type=str, dest='service_yaml', help='path to service.yaml')
  parser_become.add_argument('--type', metavar='type', type=str, dest='service_type', help='service type, needed when using --yaml')
  parser_become.add_argument('--hash', metavar='hash', type=str, dest='service_hash', help='hash of previously deployed service')
  parser_become.set_defaults(func='do_become')

  args = parser.parse_args()

  if not args.url:
    sys.stderr.write('need a url! (try --url or set GROVE_SERVER_URL)\n')
    sys.exit(1)

  if not args.name and not args.type and not args.host and not args.all and not args.online and not args.offline:
    sys.stderr.write('need some selectors! (try -a for all)\n')
    sys.exit(1)

  return args

args = init()

# Import after adjusting socket.socket
from urllib.request import urlopen, Request

def get_agents(args):
  """Returns a list of (agent_host, agent_obj) tuples. Does not return a dict, because the
  returned list is sorted."""
  agents_url = args.url.rstrip('/') + '/agents'
  response = urlopen(agents_url)
  response_obj = json.loads(response.read().decode())
  agents = []
  for agent_host, agent_obj in response_obj["agents"].items():
    matches = True
    agent_name = get_agent_name(agent_obj)
    agent_type = get_agent_type(agent_obj)
    agent_online = get_agent_online(agent_obj)

    if args.host and agent_host not in args.host:
      matches = False

    if args.name and (not agent_name or not fnmatch.fnmatch(agent_name, args.name)):
      matches = False

    if args.type and (not agent_type or not fnmatch.fnmatch(agent_type, args.type)):
      matches = False

    if args.online and not agent_online:
      matches = False

    if args.offline and agent_online:
      matches = False

    if matches:
      agents.append((agent_host, agent_obj))

  return sorted(agents, key=lambda tup: "{}/{}".format(get_agent_nametype(tup[1]), tup[0]))

def get_agent_name(agent_obj):
  if (agent_obj is not None) and (agent_obj.get("status") is not None) and (agent_obj.get("status").get("deploy") is not None):
    return agent_obj.get("status", {}).get("deploy", {}).get("name", None)
  else:
   return None

def get_agent_type(agent_obj):
  if (agent_obj is not None) and (agent_obj.get("status") is not None) and (agent_obj.get("status").get("deploy") is not None):
    return agent_obj.get("status", {}).get("deploy", {}).get("type", None)
  else:
    return None

def get_agent_nametype(agent_obj):
  agent_name = get_agent_name(agent_obj)
  agent_type = get_agent_type(agent_obj)
  if (agent_name is not None) and (agent_type is not None):
    return "{}/{}".format(agent_name, agent_type)
  else:
    return None

def get_agent_deploy_tgzs(agent_obj):
  if (agent_obj is not None) and (agent_obj.get("status") is not None) and (agent_obj.get("status").get("deploy") is not None):
    return agent_obj.get("status", {}).get("deploy", {}).get("yaml", {}).get("deploy", [])
  else:
    return None

def get_agent_online(agent_obj):
  return agent_obj.get("online", False)

def get_agent_status_time(agent_obj):
  return agent_obj.get("status_time")

def do_status(args):
  now = time.time()
  agents = get_agents(args)
  longest_name = 20
  longest_status = 0

  for agent_host, agent_obj in agents:
    agent_nametype = get_agent_nametype(agent_obj)
    if (agent_nametype is not None) and len(agent_nametype) > longest_name:
      longest_name = len(agent_nametype)
    if len(print_status(agent_obj, now)) > longest_status:
      longest_status = len(print_status(agent_obj, now))

  for agent_host, agent_obj in agents:
    agent_nametype = get_agent_nametype(agent_obj)
    agent_deploy = get_agent_deploy_tgzs(agent_obj)
    agent_online = get_agent_online(agent_obj)
    agent_status = print_status(agent_obj, now)
    if sys.stdout.isatty():
      format_str = "{:<18}{:<" + str(longest_name + 2) + "}{:<" + str(longest_status + 2) + "}{}"
    else:
      format_str = "{}\t{}\t{}\t{}"
    print(format_str.format(
      agent_host,
      agent_nametype if agent_nametype is not None else '<none>',
      agent_status,
      '; '.join(agent_deploy) if agent_deploy is not None else ''
    ))

def do_bounce(args):
  request = build_request(args, {'command': {'type': 'bounce'}})
  response_obj = post_request(args, '/do', request)
  print(json.dumps(response_obj))

def do_history(args):
  request = build_request(args, {'command': {'type': 'history'}})
  response_obj = post_request(args, '/do', request)
  print(json.dumps(response_obj))

def do_become(args):
  if (not args.service_yaml and not args.service_hash) or (args.service_yaml and (args.service_hash or not args.service_type)):
    sys.stderr.write("usage: \"become --yaml file --type type\" or \"become --hash hash\"\n")
    sys.exit(1)

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

  request = build_request(args, {
    'command': {
      'type': 'become',
      'service_yaml': service_yaml,
      'service_type': args.service_type,
      'hashed_service_parts': args.service_hash
    }
  })
  response_obj = post_request(args, '/do', request)
  print(json.dumps(response_obj))

def do_reap(args):
  request = build_request(args, {})
  response_obj = post_request(args, '/reap', request)
  print(json.dumps(response_obj))

def build_request(args, request):
  request['agents'] = []
  for agent_host, agent_obj in get_agents(args):
    request['agents'].append(agent_host)
  return request

def post_request(args, path, request):
  url = args.url.rstrip('/') + path
  response = urlopen(Request(url, json.dumps(request).encode(), {'Content-Type': 'application/json'}))
  response_obj = json.loads(response.read().decode())
  return response_obj

def print_status(agent_obj, now):
  if get_agent_online(agent_obj):
    return 'online'
  else:
    status_time = get_agent_status_time(agent_obj)
    if status_time:
      return 'offline ({})'.format(print_duration_short(now - status_time))
    else:
      return 'offline'

def print_duration_short(t):
  """Given a duration in seconds, print it using a small number of characters."""
  t = int(t)
  if t < 0:
    return "-" + print_duration_short(-1 * t)
  if t < 60:
    return str(t) + "s"
  elif t < 3600:
    return str(int(t / 60)) + "m"
  elif t < 86400:
    return str(int(t / 3600)) + "h"
  else:
    return str(int(t / 86400)) + "d"

locals().get(args.func)(args)
sys.exit(0)
