#!/usr/bin/env python
#
# Created: 20160830
# Wirawan Purwanto

"""
show-node-status.py
---------------------

Various tools to investigate node status in an SGE cluster.
This tool is a replacement and upgrade of the shell version of the tool
`node-slot-status.sh`.

Usage:

"""

import os
import re
import subprocess
import sys

#----------------------- UNDER CONSTRUCTION -----------------------
#Nothing was done yet

def node_slot_stats_raw(qstat_f, show_disabled_nodes=True):
  """Prints the node stats from `qstat -f' in raw format:
  - not printing disabled nodes
  - not showing the computational jobs that are running on these nodes
  """
  FNR = 0
  for L in qstat_f:
    FNR += 1
    FLDS = L.split()
    status_flags = FLDS[5] if (len(FLDS) > 5) else ""

    if FNR == 1 and FLDS[0] == "queuename":
      print(L)
      continue

    # Valid host status field
    if re.search(r'^[A-Za-z]', L) and len(FLDS) in (5,6) \
       and (show_disabled_nodes or ("d" not in status_flags)):
      print(L)




def help():
  msg = """\
show-node-status.py - Showing node status from SGE information

The information is mainly drawn from `qstat -f` output.

Usage: one of the following:

--raw
raw
    Shows the raw queue/node status

--stats
stats
(no argument)
    Shows the statistic summary per node type
"""


def main_default(argv, save_qstat=True):
  """Main default function:
  - By default we invoke qstat -f and prints the analysis.
  - If argv[1] is given, then we read in the file and 
    use that for the analysis.
  """
  from time import localtime, strftime

  dtime = localtime()
  dtimestr = strftime("%Y%m%d-%H%M", dtime)

  # Read the command first--what do we want to do
  if len(argv) < 2:
    cmd = "stats"
  elif argv[1] in ('--raw', 'raw'):
    cmd = "raw"
  elif argv[1] in ('--stats', 'stats'):
    cmd = "stats"
  else:
    raise ValueError, "Unknown action: "+argv[1]

  # Skip program name and first command:
  cmdargs = argv[2:]

  # Default options
  show_disabled_nodes = False

  if len(cmdargs) > 0:
    qstat_f_current = open(cmdargs[0], "r").read().splitlines()
  else:
    qstat_f_current = pipe_out(('qstat', '-f'), split=True)
    if save_qstat:
      with open("qstat-f-%s.txt" % dtimestr, "w") as F:
        F.write("\n".join(qstat_f_current))
        F.write("\n")

  if cmd == "raw":
    node_slot_stats_raw(qstat_f_current,
                        show_disabled_nodes=show_disabled_nodes,
                       )
  elif cmd == "stats":
    node_slots_stats_per_node_type(qstat_f_current,
                                   show_disabled_nodes=show_disabled_nodes,
                                  )
  else:
    raise  "Missing support for command: "+cmd



# ---------------------------------------------------------------------------
# Support tools below
# ---------------------------------------------------------------------------

def pipe_out(args, split=False, shell=False):
  """Executes a shell command, piping out the stdout to python for parsing.
  This is my customary shortcut for backtick operator.
  The result is either a single string (if split==False) or a list of strings
  with EOLs removed (if split==True)."""
  retval = subprocess.Popen(args, stdout=subprocess.PIPE, shell=shell).communicate()[0]
  if not split:
    return retval
  else:
    return retval.splitlines()


# Internal variable: don't mess!
_str_fmt_heading_rx = None
def str_fmt_heading(fmt):
  """Replaces a printf-style formatting with one suitable for table heading:
  all non-string conversions are replaced with string conversions,
  preserving the minimum widths."""
  # Originally from: $PWQMC77/scripts/cost.py and later Cr2_analysis_cbs.py .
  #
  #_str_fmt_heading_rx = None # only for development purposes
  import re
  global _str_fmt_heading_rx
  if _str_fmt_heading_rx is None:
    # Because of complicated regex, I verbosely write it out here:
    _str_fmt_heading_rx = re.compile(r"""
      (
        %                 # % sign
        (?:\([^)]+\))?    # optional '(keyname)' mapping key
        [-+#0 hlL]*       # optional conversion flag
        [0-9*]*           # optional minimum field width
      )
      ((?:\.[0-9]*)?)     # optional precision
      [^-+#*0 hlL0-9.%s]  # not conv flag, dimensions, nor literal '%',
                          # nor 's' conversion specifiers
    """, re.VERBOSE)
  return _str_fmt_heading_rx.sub(r'\1s', fmt)



# stub main code

if __name__ == "__main__" and not "get_ipython" in globals():
  main_default(sys.argv)