#!/usr/bin/env python

import json, yaml, requests, math
import eslib.debug
from eslib.esdoc import tojson
from eslib.service import status


BOOT_STATES = [status.DEAD, status.IDLE, status.PROCESSING]


def dump_formatted(format, data):
    if format == "yaml":
        yaml.safe_dump(data, sys.stdout, default_flow_style=False)
    elif format == "json":
        json.dump(data, sys.stdout)
    else:
        print >>sys.stderr, "Unrecognized output format '%s'." % format

def remote(host, verb, path, data=None, params=None):
    res = requests.request(
        verb.lower(),
        "http://%s/%s" % (host, path),
        data=tojson(data) if data else None,
        params=params,
        headers={"content-type": "application/json"},
        timeout=(3.5, 60)
    )
    if res.content:
        return json.loads(res.content)
    else:
        return None

#region Service management commands

# list [options] [--verbose] [<name_pattern>]
def cmd_list(host, format, name_patterns, rich):
    data = {
        "names": name_patterns,
    }

    res = remote(host, "get", "list", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error listing services: %s" % (error)
        sys.exit(1)

    hfmt = "%-20s %-10s %6s %-18s %11s %9s"
    ifmt = "%-20s %-10s %6s %-18s %11s %9s"
    hfmt_rich = " %6s %8s %7s %3s"
    ifmt_rich = " %6s %8s %7s %3s"

    header1 = hfmt % ("NAME", "STATUS", "PID", "HOST", "ELAPSED", "ETA")
    d = "-"
    header2 = hfmt % (d*20, d*10, d*6, d*18, d*11, d*9)

    if rich:
        header1 += hfmt_rich % ("DOCS/s", "MEMORY", "CPU", "#TH")
        header2 += hfmt_rich % (d*6, d*8, d*7, d*3)

    print header1
    print header2

    for key in sorted(list(res.keys())):
        r = res[key]
        elapsed = ""
        eta     = ""
        dps     = ""
        memory  = ""
        cpu     = ""
        threads = ""
        stats = r.get("stats")
        if stats:
            # Elapsed
            v = stats["elapsed"]
            if v is not None:
                elapsed = _hms_str(v)
            # ETA
            s = r["status"]
            if s == status.SUSPENDED:
                eta = "paused"
            elif s in [status.PROCESSING, status.STOPPING]:
                v = r["stats"]["eta"]
                if v is not None:
                    if v == -1:
                        eta = "infinite"
                    else:
                        eta = _hms_str(v)
            # Memory
            memory = _byte_str(stats["memory"])
            # CPU
            cpu = _pct_str(stats["cpu_percent"])
            # Threads
            threads = stats["threads"]
            # DPS
            dps_val = stats.get("dps")
            if dps_val is not None:
                if 0 < dps_val < 10:
                    dps = "%.2f" % dps_val
                else:
                    dps = "%.0f" % dps_val

        addr = r["host"]
        if r["port"]:
            if r["fixed_port"]:
                addr += ":%s" % str(r["port"])
            else:
                addr += ":(%s)" % str(r["port"])
        info = ifmt % (key, r["status"][:10], r["pid"], addr, elapsed, eta)

        if rich:
            info += ifmt_rich % (dps, memory, cpu, threads)

        print info

#region Value formatters

def _hm_str(value):
    if value is None: return ""
    minutes = int(math.ceil(value / 60))  # rounding minutes up
    hm = "%d:%02d" % (minutes / 60, minutes % 60)
    return hm

def _hms_str(value):
    if value is None: return ""
    seconds = int(value) % 60
    minutes = int(value) / 60
    hms = "%d:%02d:%02d" % (minutes / 60, minutes % 60, seconds)
    return hms

def _pct_str(value):
    if value is None: return ""
    return "%0.1f%%" % value

def _byte_str(value):
    if value is None: return ""
    return eslib.debug.byte_size_string(value, 1)

#endregion Value formatters

# stats [options] [<name>] [<stat_fields...>]
def cmd_stats(host, format, name, stat_fields):
    data = {
        "ids"   : [name] if name else None
    }
    res = remote(host, "get", "stats", data)
    error = res.get("error")
    if format:
        if error:
            return dump_formatted(format, {"error": error})
        else:
            if stat_fields:
                ret = {}
                info = res.get(name) if name else res.itervalues().next()
                for key in sorted(info.keys()):
                    if key == "stats":
                        continue
                    if key in stat_fields:
                        ret[key] = info[key]
                stats = info.get("stats")
                if stats:
                    for key in stats.keys():
                        if key in stat_fields:
                            ret[key] = stats[key]
                return dump_formatted(format, ret)
            else:
                return dump_formatted(format, res)
    error = res.get("error")
    if not error and res.get(name):
        error = res[name].get("error")
    if error:
        print >>sys.stderr, "Error getting stats for service '%s': %s" % (name, error)
        sys.exit(1)

    indent = "    " if not stat_fields else ""

    info = res.get(name) if name else res.itervalues().next()
    if info:
        if not stat_fields:
            print "INFO FOR SERVICE '%s':" % name
        for key in sorted(info.keys()):
            if key == "stats":
                continue
            if stat_fields and not key in stat_fields:
                continue
            value = info[key]
            print "%s%-15s = %s" % (indent, key, value)
        stats = info.get("stats")
        if stats:
            if not stat_fields:
                print "STATS:"
            for key in sorted(stats.keys()):
                if stat_fields and not key in stat_fields:
                    continue
                value = stats[key]
                if key in ["uptime", "elapsed", "eta"]:
                    value = _hms_str(value)
                elif key in ["cpu_percent", "cpu_percent_max"]:
                    value = _pct_str(value)
                elif key in ["memory", "memory_max"]:
                    value = _byte_str(value)
                elif key in ["dps"]:
                    value = "" if value is None else ("%.2f" % value)
                print "%s%-15s = %s" % (indent, key, value)

# add [options] [--start] <name> [-c <config_tag>] [-s <server>]
def cmd_add(host, format, name, boot_state, config_tag, server_address, auto_start):
    if boot_state and not boot_state.lower() in [b.lower() for b in BOOT_STATES]:
        print >>sys.stderr, "Invalid boot state '%s'. Must be one of [%s]" % (boot_state, ", ".join([("'%s'" % b) for b in BOOT_STATES]))
        sys.exit(3)

    remote_host = server_address or "localhost"
    remote_port = None
    addr = remote_host.split(":")
    if len(addr) == 2:
        remote_host = addr[0]
        remote_port = addr[1]

    data = {
        "id"         : name,
        "host"       : remote_host,
        "port"       : remote_port,
        "boot_state" : boot_state,
        "config_key" : config_tag,
        "start"      : auto_start,
    }
    res = remote(host, "put", "add", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error adding service '%s': %s" % (name, error)
        sys.exit(1)

# remove [options] [--all] [--stop] <name>
def cmd_remove(host, format, names, all, auto_stop):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
        "stop"  : auto_stop
    }
    res = remote(host, "delete", "remove", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error removing service(s) %s: %s" % (names, error)

# run [options] [-all] [--start] <names...>
def cmd_run(host, format, names, all, start):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
        "start" : start,
    }
    res = remote(host, "post", "run", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error launching service(s) %s: %s" % (names, error)

# shutdown [options] [--all] [--wait] <names...>
def cmd_shutdown(host, format, names, all, wait):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
        "wait"  : wait,
    }
    res = remote(host, "delete", "shutdown", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error shutting down service(s) %s: %s" % (names, error)

# kill [options] [--all] [--force] <names...>
def cmd_kill(host, format, names, all, force):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
        "force" : force
    }
    res = remote(host, "delete", "kill", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error killing processes for service(s) %s: %s" % (names, error)
        sys.exit(1)

# reload [options] [-all] <names...>
def cmd_reload(host, format, names, all):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
    }
    res = remote(host, "post", "reboot", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error reloading service(s) %s: %s" % (names, error)
        sys.exit(1)

# start [options] [--all] <names...>
def cmd_start(host, format, names, all):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
    }
    res = remote(host, "post", "processing_start", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error starting processing for service(s) %s: %s" % (names, error)
        sys.exit(1)

# stop [options] [--all] [--wait] <names...>
def cmd_stop(host, format, names, all, wait):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
        "wait"  : wait
    }
    res = remote(host, "post", "processing_stop", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error stopping processing for service(s) %s: %s" % (names, error)
        sys.exit(1)

# abort [options] [--all] <names...>
def cmd_abort(host, format, names, all):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
    }
    res = remote(host, "post", "processing_abort", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error stopping processing for service(s) %s: %s" % (names, error)
        sys.exit(1)

# suspend [options] [--all] <names...>
def cmd_suspend(host, format, names, all):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
    }
    res = remote(host, "post", "processing_suspend", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error suspending processing for service(s) %s: %s" % (names, error)
        sys.exit(1)

# resume [options] [--all] <names...>
def cmd_resume(host, format, names, all):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)

    data = {
        "ids"   : names,
        "all"   : all,
    }
    res = remote(host, "post", "processing_resume", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error resuming processing for service(s) %s: %s" % (names, error)
        sys.exit(1)

# boot [options] [--all] <names...>
def cmd_boot(host, format, names, all):
    # Default to all if no service names are given
    if not names:
        all = True

    data = {
        "ids"   : names,
        "all"   : all,
    }
    res = remote(host, "post", "boot", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error booting service(s) %s: %s" % (names, error)

# set-boot [options] <boot-state> [--all] <names...>
def cmd_set_boot(host, format, boot_state, names, all):
    if not names and not all:
        print >>sys.stderr, "Missing service name(s) or '--all'."
        sys.exit(2)
    if boot_state is None or not boot_state.lower() in [b.lower() for b in BOOT_STATES]:
        print >>sys.stderr, "Invalid boot state '%s'. Must be one of [%s]" % (boot_state, ", ".join([("'%s'" % b) for b in BOOT_STATES]))
        sys.exit(3)

    data = {
        "ids"        : names,
        "boot_state" : boot_state,
        "all"        : all,
    }
    res = remote(host, "post", "set_boot_state", data)
    if format: return dump_formatted(format, res)
    error = res.get("error")
    if error:
        print >>sys.stderr, "Error setting boot state for service(s) %s: %s" % (names, error)
        sys.exit(1)

#endregion Service management commands


import argparse, sys, os

def main():
    usage = \
"""
Usage:
  %(prog)s list     [options] [--verbose] [<name_patterns...>]
  %(prog)s stats    [options] <name> [<stat_fields...>]

  %(prog)s add      [options] [--start] <name>
  %(pspc)s          [-b <boot_state>] -c <config_tag>] [-s <server_address>]
  %(prog)s remove   [options] [--stop] <names...>

  %(prog)s run      [options] [--all] [--start] <names...>
  %(prog)s shutdown [options] [--all] [--wait]  <names...>
  %(prog)s kill     [options] [--all] [--force] <names...>

  %(prog)s start    [options] [--all] <names...>
  %(prog)s stop     [options] [--all] [--wait] <names...>
  %(prog)s abort    [options] [--all] <names...>
  %(prog)s suspend  [options] [--all] <names...>
  %(prog)s resume   [options] [--all] <names...>

  %(prog)s set-boot [options] <state> [--all] <names...>
  %(prog)s boot     [options] [--all] <names...>
  %(prog)s reboot   [options] [--all] <names...>

Common options:
  --host=<host>
  --json
  --yaml

Boot states are one of %(boot_states)s.
"""

    # Set up parser

    parser_desc = "Service management client for document processing services."
    parser = argparse.ArgumentParser(description=parser_desc)
    parser._actions[0].help = argparse.SUPPRESS
    subparsers = parser.add_subparsers(
        title = "subcommands",
        #description = "valid subcommands",
        help = "additional help",
        dest = "command")  # Name of variable to host which command was specified

    parser_common = argparse.ArgumentParser(add_help=False)
    parser_common.add_argument("--host", type=str, required=False)
    parser_common.add_argument("--yaml", action="store_true")
    parser_common.add_argument("--json", action="store_true")

    # nargs="?" => 0 or 1 (NOT A LIST; either value or None... this is the weird shit)
    # nargs="*" => 0 or more (list)
    # nargs="+" => 1 or more (list)
    # nargs="2" => 2 (list)
    # nargs="1" => 1 (list)

    # list [options] [--verbose] [<name_patterns...>]
    desc = "List services."
    parser_list = subparsers.add_parser("list", description=desc, parents=[parser_common])
    parser_list._actions[0].help = argparse.SUPPRESS
    parser_list.add_argument("names", type=str, nargs="*")
    parser_list.add_argument("-v", "--verbose", "--rich", dest="rich", action="store_true")
    # stats [options] [<name>] [<stat_fields...>]
    desc = "Display runtime statistics for service."
    parser_stats = subparsers.add_parser("stats", description=desc, parents=[parser_common])
    parser_stats._actions[0].help = argparse.SUPPRESS
    parser_stats.add_argument("name", type=str, nargs="?")
    parser_stats.add_argument("fields", type=str, nargs="*")  # * = 0 or more

    # add [options] [--start] <name> [-c <config_tag>] [-s <server_address>]
    desc = "Register a service that can be ran. Registration is stored."
    parser_add = subparsers.add_parser("add", description=desc, parents=[parser_common])
    parser_add._actions[0].help = argparse.SUPPRESS
    parser_add.add_argument("name", type=str, nargs=1)
    parser_add.add_argument("-c", "--tag", dest="config_tag", type=str, required=False)
    parser_add.add_argument("-s", "--server", type=str, required=False)
    parser_add.add_argument("-b", "--boot-state", dest="boot_state", type=str, required=False)
    parser_add.add_argument("--all", action="store_true")
    parser_add.add_argument("--start", dest="auto_start", action="store_true")
    # remove [options] [--all] [--stop] <names...>
    desc = "Unregister a service."
    parser_remove = subparsers.add_parser("remove", description=desc, parents=[parser_common])
    parser_remove._actions[0].help = argparse.SUPPRESS
    parser_remove.add_argument("names", type=str, nargs="+")
    parser_remove.add_argument("--all", action="store_true")
    parser_remove.add_argument("--stop", dest="auto_stop", action="store_true")

    # run [options] [--all] [--start] <names...>
    desc = "Start the registered service as an OS process and run the service. (Not !document processing!)"
    parser_run = subparsers.add_parser("run", description=desc, parents=[parser_common])
    parser_run._actions[0].help = argparse.SUPPRESS
    parser_run.add_argument("names", type=str, nargs="*")
    parser_run.add_argument("--all", action="store_true")
    parser_run.add_argument("--start", dest="auto_start", action="store_true")
    # shutdown [options] [--all] [--wait] <names...>
    desc = "Stop processing and shut down service."
    parser_shutdown = subparsers.add_parser("shutdown", description=desc, parents=[parser_common])
    parser_shutdown._actions[0].help = argparse.SUPPRESS
    parser_shutdown.add_argument("names", type=str, nargs="*")
    parser_shutdown.add_argument("--wait", action="store_true")
    parser_shutdown.add_argument("--all", action="store_true")
    # kill [options] [--all] [--force] <names...>
    desc = "Kill OS process. Meant for cleaning up dead/hanging processes."
    parser_kill = subparsers.add_parser("kill", description=desc, parents=[parser_common])
    parser_kill._actions[0].help = argparse.SUPPRESS
    parser_kill.add_argument("names", type=str, nargs="*")
    parser_kill.add_argument("--force", action="store_true")
    parser_kill.add_argument("--all", action="store_true")

    # start [options] [--all] <names...>
    desc = "Start document processing."
    parser_start = subparsers.add_parser("start", description=desc, parents=[parser_common])
    parser_start._actions[0].help = argparse.SUPPRESS
    parser_start.add_argument("names", type=str, nargs="*")
    parser_start.add_argument("--all", action="store_true")
    # stop [options] [--all] [--wait] <names...>
    desc = "Stop document processing."
    parser_stop = subparsers.add_parser("stop", description=desc, parents=[parser_common])
    parser_stop._actions[0].help = argparse.SUPPRESS
    parser_stop.add_argument("names", type=str, nargs="*")
    parser_stop.add_argument("--wait", action="store_true")
    parser_stop.add_argument("--all", action="store_true")
    # abort [options] [--all] <names...>
    desc = "Abort document processing."
    parser_abort = subparsers.add_parser("abort", description=desc, parents=[parser_common])
    parser_abort._actions[0].help = argparse.SUPPRESS
    parser_abort.add_argument("names", type=str, nargs="*")
    parser_abort.add_argument("--all", action="store_true")
    # suspend [options] [--all] <names...>
    desc = "Suspend document processing."
    parser_suspend = subparsers.add_parser("suspend", description=desc, parents=[parser_common])
    parser_suspend._actions[0].help = argparse.SUPPRESS
    parser_suspend.add_argument("names", type=str, nargs="*")
    parser_suspend.add_argument("--all", action="store_true")
    # resume [options] [--all] <names...>
    desc = "Resume document processing."
    parser_resume = subparsers.add_parser("resume", description=desc, parents=[parser_common])
    parser_resume._actions[0].help = argparse.SUPPRESS
    parser_resume.add_argument("names", type=str, nargs="*")
    parser_resume.add_argument("--all", action="store_true")

    # set-boot [options] <boot-state> [--all] <names...>
    desc = "Set boot state; one of 'dead', 'idle', 'processing'."
    parser_setboot = subparsers.add_parser("set-boot", description=desc, parents=[parser_common])
    parser_setboot._actions[0].help = argparse.SUPPRESS
    parser_setboot.add_argument("boot_state", type=str, nargs=1)
    parser_setboot.add_argument("names", type=str, nargs="*")
    parser_setboot.add_argument("--all", action="store_true")
    # boot [options] [--all] <names...>
    desc = "Spin up services with boot state and make sure they are running or processing according to their boot states."
    parser_boot = subparsers.add_parser("boot", description=desc, parents=[parser_common])
    parser_boot._actions[0].help = argparse.SUPPRESS
    parser_boot.add_argument("names", type=str, nargs="*")
    parser_boot.add_argument("--all", action="store_true")
    # reboot [options] [--all] <names...>
    desc = "Reboot services. It will launch with new executable code and attempt run state according to boot-state settings."
    parser_reboot = subparsers.add_parser("reboot", description=desc, parents=[parser_common])
    parser_reboot._actions[0].help = argparse.SUPPRESS
    parser_reboot.add_argument("names", type=str, nargs="*")
    parser_reboot.add_argument("--all", action="store_true")

    # Parse

    if len(sys.argv) == 1:
        print usage.strip() % {
            "prog": parser.prog,
            "pspc": " "*len(parser.prog),
            "boot_states" :  ", ".join([("'%s'" % b) for b in BOOT_STATES])
        }
        sys.exit(1)

    args = parser.parse_args()

    # Prepare common vars from parsing

    format = None
    if args.yaml:
        format = "yaml"
    elif args.json:
        format = "json"
    host = args.host or os.environ.get("ESLIB_SERVICE_MANAGER") or "localhost:5000"

    # Call commands according to parsed verbs

    if   args.command == "list"      : cmd_list    (host, format, args.names, args.rich)
    elif args.command == "stats"     : cmd_stats   (host, format, args.name if args.name else None, args.fields)
    elif args.command == "add"       : cmd_add     (host, format, args.name[0], args.boot_state, args.config_tag, args.server, args.auto_start)
    elif args.command == "remove"    : cmd_remove  (host, format, args.names, args.all, args.auto_stop)
    elif args.command == "run"       : cmd_run     (host, format, args.names, args.all, args.auto_start)
    elif args.command == "shutdown"  : cmd_shutdown(host, format, args.names, args.all, args.wait)
    elif args.command == "kill"      : cmd_kill    (host, format, args.names, args.all, args.force)
    elif args.command == "start"     : cmd_start   (host, format, args.names, args.all)
    elif args.command == "stop"      : cmd_stop    (host, format, args.names, args.all, args.wait)
    elif args.command == "abort"     : cmd_abort   (host, format, args.names, args.all)
    elif args.command == "suspend"   : cmd_suspend (host, format, args.names, args.all)
    elif args.command == "resume"    : cmd_resume  (host, format, args.names, args.all)
    elif args.command == "set-boot"  : cmd_set_boot(host, format, args.boot_state[0], args.names, args.all)
    elif args.command == "boot"      : cmd_boot    (host, format, args.names, args.all)
    elif args.command == "reboot"    : cmd_reload  (host, format, args.names, args.all)

if __name__ == '__main__':
    main()
