#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Wooki job info.

Shows job information, like qstat, but customised for jobs on the local
cluster. Shows all user jobs by default, use the '--user' argument to
specify a user or list of users to show jobs for.

"""

from __future__ import print_function, division
import argparse
import re
import signal
import sys
from datetime import datetime
from getpass import getuser
from os.path import dirname, realpath

from dateutil.parser import parse

# Put the parent directory first in the path
sys.path.insert(1, dirname(dirname(realpath(__file__))))
from wsge.qstat import qstat, qstatf

# If output goes into head or less without finishing, it will terminate with
# IOError, this makes pipe errors get ignored.
signal.signal(signal.SIGPIPE, signal.SIG_DFL)


def commandline():
    """Process commandline options in here. Return the argparse object."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("-u", "--user", nargs='+',
                        help="User or list of users to show jobs for. Names "
                        "may use regular expressions for wildcard matches.")
    parser.add_argument("-q", "--queue", action='store_true',
                        help="Show queue name in addition to the node name.")
    parser.add_argument("-f", "--full", action='store_true',
                        help="Show full output by nodes, like 'qstat -f'.")
    parser.add_argument("-n", "--no-colour", "--no-color", dest='colour',
                        action="store_false", help="Turn off coloured output.")
    options = parser.parse_args()

    return options


class ColourDict(dict):
    """
    Store items as colourised versions of themselves, changing the colour
    for each new item.
    """
    def __init__(self, width=12, mapping=None):
        self.width = width
        super(ColourDict, self).__init__(self)
        if mapping is not None:
            self.update(mapping)

    def __missing__(self, key):
        if not hasattr(self, 'counter'):
            self.counter = 1
        self[key] = "\033[38;5;{colour_id}m{key:<{width}}\033[0m".format(
            colour_id=self.counter, key=key, width=self.width)
        self.counter += 1
        return self[key]


class AutoDict(dict):
    """
    Add the missing entries as themselves to a dictionary.
    """
    def __init__(self, width=12):
        self.width = width
        super(AutoDict, self).__init__(self)

    def __missing__(self, key):
        self[key] = "{key:<{width}}".format(key=key, width=self.width)
        return self[key]


def parse_number(number):
    """
    Grok the value of the number from a string, deals with the M, G that appear
    in sge.

    """
    suffixes = {
        'K': 1024,
        'M': 1024**2,
        'G': 1024**3,
        'T': 1024**4}

    if isinstance(number, (int, float)):
        return number

    if '.' in number:
        fstr = re.search(r'-?(\d+(\.\d*)?|\.\d+)([eE][+-]?\d+)?', number)
        value = float(fstr.group())
    else:
        istr = re.search(r'-?\d+', number)
        value = int(istr.group())

    for suffix, multiplier in suffixes.items():
        if suffix in number:
            value *= multiplier

    return value


def progressbar(value, maximum, width=12, show_value=True, colour=True,
                label=None):
    """
    Generate a visual representation of progress.
    """
    bar = []
    if label is not None:
        bar.append(label)
    pvalue = parse_number(value)
    pmaximum = parse_number(maximum)

    used = int(width*pvalue/pmaximum)
    if used > pmaximum:
        used = width
    elif used < 0:
        used = 0
    unused = width - used
    bar.extend(['[', '|'*used, '.'*unused, ']'])

    if show_value:
        bar.append(' {value}/{max}'.format(value=value, max=maximum))

    return "".join(bar)


def numeric_sort(name):
    """
    Pull all digits from the name and make a list with all number in it.
    """
    return [int(x) for x in re.findall(r'[0-9]+', name)]


def qstat_priority(options):
    """Process qstat and print formatted output to terminal."""

    # Get data
    jobs = qstat()

    # Use regular expressions to prune the job list
    if options.user is not None:
        user_regexes = [re.compile('^{0}$'.format(user_regex))
                        for user_regex in options.user]
        jobs = [job for job in jobs if
                any(user_regex.match(job['JB_owner'])
                    for user_regex in user_regexes)]

    # are there job arrays?
    task_width = max([len(job.get('tasks', '')) for job in jobs] + [0])
    if task_width > 0:
        task_width += 1
        task_str = ' tasks'
    else:
        task_str = ''

    # how wide? include 8 for empty case so header looks reasonable
    name_width = max([8] + [len(job['JB_name']) for job in jobs])

    # header              "xky        "      job age             t
    print("job-ID   prior      user          state  run/wait time       queue"
          "            slots  {1:{0}}{2}"
          "".format(name_width, 'name', task_str))
    print("—"*(85+name_width+task_width))

    # Get rid of the microseconds so that the output is only reported to the
    # second
    now = datetime.now().replace(microsecond=0)

    # Track if we are in pending jobs to add a separator
    pending_break = False
    for job in jobs:
        if not pending_break and 'q' in job['state']:
            print("—"*(85+name_width+task_width))
            pending_break = True

        # Parse the dates into deltas -> more useful?
        try:
            runtime = now - parse(job['JAT_start_time'])
        except KeyError:
            runtime = now - parse(job['JB_submission_time'])

        # non running jobs have no queue
        if job['queue_name'] is not None:
            if options.queue:
                queue = job['queue_name']
            else:
                queue = job['queue_name'].split('@')[1].split('.')[0]
        else:
            queue = ''

        # job_array?
        if 'tasks' in job:
            tasks = ' ' + job['tasks']
        else:
            tasks = ''

        print("{0[JB_job_number]:<7s}  {0[JAT_prio]:<9s}  {0[JB_owner]:<12}  "
              "{0[state]:<5s}  {runtime:>18s}  {queue:<15}  {0[slots]:>5}  "
              "{0[JB_name]:<{name_width}s}{tasks:<s}"
              "".format(job, runtime=runtime, queue=queue,
                        name_width=name_width, tasks=tasks))

    return


def qstat_full(options):
    """
    Process jobs and arrange like 'qstst -f', but nicer.
    """

    all_queues, users, jobs = qstatf()

    # Use regular expressions to prune the job list
    if options.user is not None:
        user_regexes = [re.compile('^{0}$'.format(user_regex))
                        for user_regex in options.user]
        jobs = [job for job in jobs if
                any(user_regex.match(job['JB_owner'])
                    for user_regex in user_regexes)]

    # are there job arrays?
    task_width = max([len(job.get('tasks', '')) for job in jobs] + [0])
    if task_width > 0:
        task_width += 1
        task_str = ' tasks'
    else:
        task_str = ''

    # how wide? include 8 for empty case so header looks reasonable
    name_width = max([8] + [len(job['JB_name']) for job in jobs])

    # header
    print("job-ID   prior      user          state  run/wait time       queue"
          "            slots  {1:{0}}{2}"
          "".format(name_width, 'name', task_str))

    # Get rid of the microseconds so that the output is only reported to the
    # second
    now = datetime.now().replace(microsecond=0)

    # Setup how people's jobs look:
    if options.colour:
        this_user = {getuser():
                     "\033[104;39;1m{user:<12}\033[0m".format(user=getuser())}
        usercodes = ColourDict(width=12, mapping=this_user)
        queues = ColourDict(width=15)
    else:
        usercodes = AutoDict(width=12)
        queues = AutoDict(width=15)

    jobfmt = ("{0[JB_job_number]:<7s}  {0[JAT_prio]:<9s}  {user}  "
              "{0[state]:<5s}  {runtime:>18s}  {queue}  {0[slots]:>5}  "
              "{0[JB_name]:<{name_width}s}{tasks:<s}")

    for node in sorted(all_queues, key=numeric_sort):
        info = all_queues[node]
        print("—"*(85+name_width+task_width))
        #print(all_queues[node]['resource'])

        nodestr = ["\033[1m>> {node:10} ".format(node=node, nodeinfo=info),
                   progressbar(info['slots_used'],
                               info['slots_total'],
                               width=int(info['slots_total']),
                               label='Slots:  ')]
        if 'load_avg' in info['resource']:
            nodestr.append('Load: {0:.2f}'.format(
                float(info['resource']['load_avg'])))
        if 'mem_total' in info['resource']:
            nodestr.append(progressbar(
                info['resource'].get('mem_used', '0.0 G'),
                info['resource']['mem_total'],
                width=20, label="\n>>              Memory: "))
        print("  ".join(nodestr) + "\033[38;0m")

        for job in jobs:
            if 'node' in job and job['node'] == node:
                user = usercodes[job['JB_owner']]
                runtime = now - parse(job['JAT_start_time'])
                queue = queues[job['JB_queue'].split('@')[0]]
                if 'tasks' in job:
                    tasks = ' ' + job['tasks']
                else:
                    tasks = ''
                print(jobfmt.format(job, user=user, runtime=runtime,
                                    queue=queue, tasks=tasks,
                                    name_width=name_width))

    print("—"*(85+name_width+task_width))
    print(" Pending jobs ")
    print("—"*(85+name_width+task_width))
    for job in jobs:
        if 'q' in job['state']:
            user = usercodes[job['JB_owner']]
            runtime = now - parse(job['JB_submission_time'])
            queue = ' '*15  # Expects 15 wide (ansi codes excepted)
            if 'tasks' in job:
                tasks = ' ' + job['tasks']
            else:
                tasks = ''
            print(jobfmt.format(job, user=user, runtime=runtime, queue=queue,
                                tasks=tasks, name_width=name_width))

    return


def main():
    """Process options and call the qstat formatting code."""
    options = commandline()

    if options.full:
        qstat_full(options)
    else:
        qstat_priority(options)


if __name__ == '__main__':
    main()
