# -*- coding: iso-8859-1 -*-

# ==============================================================================
# COPYRIGHT (C) 1991 - 2003  EDF R&D                  WWW.CODE-ASTER.ORG
# THIS PROGRAM IS FREE SOFTWARE; YOU CAN REDISTRIBUTE IT AND/OR MODIFY
# IT UNDER THE TERMS OF THE GNU GENERAL PUBLIC LICENSE AS PUBLISHED BY
# THE FREE SOFTWARE FOUNDATION; EITHER VERSION 2 OF THE LICENSE, OR
# (AT YOUR OPTION) ANY LATER VERSION.
#
# THIS PROGRAM IS DISTRIBUTED IN THE HOPE THAT IT WILL BE USEFUL, BUT
# WITHOUT ANY WARRANTY; WITHOUT EVEN THE IMPLIED WARRANTY OF
# MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. SEE THE GNU
# GENERAL PUBLIC LICENSE FOR MORE DETAILS.
#
# YOU SHOULD HAVE RECEIVED A COPY OF THE GNU GENERAL PUBLIC LICENSE
# ALONG WITH THIS PROGRAM; IF NOT, WRITE TO EDF R&D CODE_ASTER,
#    1 AVENUE DU GENERAL DE GAULLE, 92141 CLAMART CEDEX, FRANCE.
# ==============================================================================

"""
This module gives functions to manipulate Code_Aster jobs :
- return informations from astk server configuration,
- edit output or error file,
- get job status,
- kill and delete the job files,
- search strings in output of a job,
- purge 'flasheur' directory.
The functions are called by an ASTER_RUN object.
"""

__revision__ = "$Id: as_job.py 2627 2006-12-18 07:50:58Z courtois $"


# ----- differ messages translation
def _(mesg):
   return mesg

import os
import signal
import glob
import re
import getpass

from as_profil import ASTER_PROFIL

#-------------------------------------------------------------------------------
def SetParser(run):
   """Configure the command-line parser, add options name to store to the list,
   set actions informations.
   run : ASTER_RUN object which manages the execution
   """
   acts_descr = {
      'info' : {
         'method' : Info,
         'syntax' : '',
         'help'   : 'returns informations from astk server configuration : '\
                    'batch, interactive (yes/no, limits), compute nodes, '\
                    'versions'
      },
      'actu' : {
         'method' : Actu,
         'syntax' : 'job_number job_name mode',
         'help'   : 'returns the state, diagnosis, execution node, spent '\
                    'cpu time and working directory of a job'
      },
      'edit' : {
         'method' : Edit,
         'syntax' : 'job_number job_name mode output|error DISPLAY',
         'help'   : 'opens output or error file on the provided display'
      },
      'del' : {
         'method' : Del,
         'syntax' : 'job_number job_name mode [node] [--signal=...]',
         'help'   : 'kill a job and delete related files'
      },
      'purge_flash' : {
         'method' : Purge,
         'syntax' : 'job_number1 [job_number2 [...]]]',
         'help'   : 'delete files of jobs which are NOT in the list'
      },
      'tail' : {
         'method' : Tail,
         'syntax' : 'job_number job_name mode fdest nb_lines [regexp]',
         'help'   : 'output the last part of fort.6 file or filter lines ' \
                    'matching a pattern'
      },
   }
   opts_descr = {
      'signal' : {
         'args'   : ('--signal', ),
         'kwargs' : {
            'action'  : 'store',
            'default' : 'KILL',
            'choices' : ('KILL', 'USR1'),
            'dest'    : 'signal',
            'help'    : 'signal to the job (KILL|USR1)'
         }
      },
   }
   title = 'Options for operations on jobs'
   run.SetActions(
         actions_descr = acts_descr,
         actions_order = ['info', 'actu', 'edit', 'tail', 'del', 'purge_flash'],
         group_options=True, group_title=title, actions_group_title=False,
         options_descr = opts_descr,
   )

#-------------------------------------------------------------------------------
def Info(run, *args):
   """Return informations from astk server configuration.
   """
   if len(args)>0:
      run.parser.error(_("'--%s' requires no argument") % run.current_action)
   
   on_machref = run.get('rep_agla', 'local') != 'local'
   
   # astk server version
   l_v = run.__version__.split('.')
   try:
      svers = '.'.join(['%02d' % int(i) for i in l_v])
   except ValueError:
      svers = run.__version__
   output = ["@SERV_VERS@", svers, "@FINSERV_VERS@"]

   # all these parameters are necessary   
   output.append("@PARAM@")
   output.extend(['%s : %s' % (p, run[p]) for p in ('plate-forme', 'batch',
      'batch_memmax', 'batch_tpsmax', 'batch_nbpmax', 'interactif',
      'interactif_memmax', 'interactif_tpsmax', 'interactif_nbpmax')])
   output.append("@FINPARAM@")

   # Code_Aster versions
   output.append("@VERSIONS@")
   output.extend(['vers : %s' % v for v in run.get('vers', '').split()])
   output.append("@FINVERSIONS@")

   output.append("@NOEUDS@")
   output.extend(['noeud : %s' % n for n in run.get('noeud', '').split()])
   output.append("@FINNOEUDS@")
   
   if on_machref:
      output.extend(["@MACHREF@", "@FINMACHREF@"])
      output.extend(["@REX_URL@", run.get('rex_url', ''), "@FINREX_URL@"])
      output.extend(["@REX_REPFICH@", run.get('rep_rex', ''), "@FINREX_REPFICH@"])
      output.extend(["@MAIL_ATA@", run.get('mail_ata', ''), "@FINMAIL_ATA@"])
      
      # read identAster file
      f_ident = os.path.join(run['rep_identAster'], 'identAster')
      if not os.path.isfile(f_ident):
         run.Sortie(4)
      f = open(f_ident, 'r')
      l_ident = [line.strip().replace('\t', ' ') \
               for line in f if re.search('^[ ]*%', line)==None]
      f.close()
      d_user = {}
      for l in l_ident:
         l_i = l.split()
         if len(l_i) < 2:
            continue
         d_user[l_i[1]] = l_i
      
      # extract user info
      role = 'UA'
      details = ''
      user = getpass.getuser()
      if d_user.has_key(user):
         l = d_user[user]
         if len(l) >= 6:
            role = l[5].replace('/', ' ')
            details = ('%s %s %s %s' % (l[0], l[1], l[2], l[3])).replace('/', ' ')
      
      output.extend(["@ROLE@", role, "@FINROLE@"])
      if details:
         output.extend(["@INFOID@", details, "@FININFOID@"])

   # message of the day
   output.append("@MOTD@")
   if run.get('motd', '') != '' and os.path.exists(run['motd']):
      output.append(open(run['motd'], 'r').read())
   output.append("@FINMOTD@")
   
   print os.linesep.join(output)

#-------------------------------------------------------------------------------
def _nodepara(run, nodename, para):
   """Return of `para` from config file or config_`nodename` file if exists.
   """
   val = run[para]
   fcn = os.path.join(run.confdir, 'config_%s' % nodename)
   if os.path.isfile(fcn):
      mdict = {}
      run._read_rc(fcn, mdict)
      val = mdict.get(para, val)
   return val

#-------------------------------------------------------------------------------
def _sumtps(l_scpu):
   """Return the sum in seconds of the times given as strings in `l_scpu`.
   Format of each times in [dd:[hh:[mm:[ss]]]]
   """
   tot = 0
   if len(l_scpu) < 1:
      tot = '_'
   for scpu in l_scpu:
      try:
         s = [int(s) for s in scpu.split(':')]
      except ValueError:
         s = []
      if   len(s) == 4:
         tot += 86400*s[0]+3600*s[1]+60*s[2]+s[3]
      elif len(s) == 3:
         tot += 3600*s[0]+60*s[1]+s[2]
      elif len(s) == 2:
         tot += 60*s[0]+s[1]
      elif len(s) == 1:
         tot += s[0]
   return tot

#-------------------------------------------------------------------------------
def Func_actu(run, *args):
   """Return state, diagnosis, node, cpu time and working directory of a job.
   """
   if len(args) != 3:
      run.parser.error(_("'--%s' takes exactly %d arguments (%d given)") % \
         (run.current_action, 3, len(args)))
   
   njob, nomjob, mode = args
   # defaults
   etat  = '_'
   diag  = '_'
   node  = '_'
   tcpu  = '_'
   wrk   = '_'
   queue = '_'
   psout = ''
   on_machref = run.get('rep_agla', 'local') != 'local'
   # the real job id may differ
   jobid = str(njob)
   
   # 1. get information about the job
   # 1.1. batch mode
   if mode == "batch":
      m = 'b'
      # for LSF
      if run['batch_nom'] == 'LSF':
         cmd = ''
         if run.get('batch_ini', '') != '':
            cmd = '. %s ; ' % run['batch_ini']
         cmd += '%s -uall' % run['batch_job']
         jret, out = run.Shell(cmd)
         mat = re.search('(^ *%s.*)' % njob, out, re.MULTILINE)
         if mat == None:
            etat = "ENDED"
         else:
            lin = mat.group(1).split()
            if len(lin) >= 3:
               etat = lin[2]
               if etat != "PEND":
                  node = lin[5]
            if etat.find('SUSP')>-1:
               etat = 'SUSPENDED'
      # for PBS
      elif run['batch_nom'] == 'PBS':
         cmd = ''
         if run.get('batch_ini', '') != '':
            cmd = '. %s ; ' % run['batch_ini']
         cmd += '%s -f %s' % (run['batch_job'], njob)
         jret, out = run.Shell(cmd)
         # real job id
         mjid = re.search('Job Id: *(\S+)', out, re.MULTILINE)
         if mjid != None:
            jobid = mjid.group(1)
         # job state
         metat = re.search(' job_state *= *(\S+)', out, re.MULTILINE)
         if metat == None:
            etat = "ENDED"
         else:
            etat = { 'R' : 'RUN', 'E'  :'RUN',
                     'W' : 'PEND', 'T' : 'PEND', 'Q' : 'PEND',
                     'S' : 'SUSPENDED', 'H' : 'SUSPENDED',
                   }.get(metat.group(1), '?')
         # queue
         mqueue = re.search(' queue *= *(\S+)', out, re.MULTILINE)
         if mqueue != None:
            queue = mqueue.group(1)
         # exec host
         mnode = re.search(' exec_host *= *(\S+)/', out, re.MULTILINE)
         if mnode != None:
            node = mnode.group(1)
      # for Sun Grid Engine
      elif run['batch_nom'] == 'SunGE':
         cmd = ''
         if run.get('batch_ini', '') != '':
            cmd = '. %s ; ' % run['batch_ini']
         cmd += run['batch_job']
         jret, out = run.Shell(cmd)
         mat = re.search('(^ *%s.*)' % njob, out, re.MULTILINE)
         if mat == None:
            etat = "ENDED"
         else:
            lin = mat.group(1).split()
            if len(lin) >= 5:
               etat = lin[4]
               if re.search('[wh]+', etat) != None:
                  etat = "PEND"
               elif re.search('[sST]+', etat) != None:
                  etat = "SUSPENDED"
               else:
                  etat = "RUN"
               if len(lin) >= 8:
                  l_q = lin[7].split('@')
                  queue = l_q[0]
                  if len(l_q) > 1:
                     node  = l_q[1].split('.')[0]
            # because I don't known the running node on which to call 'ps'
            if etat == 'RUN' and node in ('', '_'):
               cmd = ''
               if run.get('batch_ini', '') != '':
                  cmd = '. %s ; ' % run['batch_ini']
               cmd += '%s -j %s' % (run['batch_job'], njob)
               jret, out = run.Shell(cmd)
               expr = re.compile('^usage *[0-9]* *: *cpu=([0-9:]+)', re.MULTILINE)
               l_field = expr.findall(out)
               tcpu = _sumtps(l_field)
      else:
         run.Mess(_('unknown batch scheduler'), _('<F>_ERROR'))

   # 1.2. interactive mode
   elif mode == "interactif":
      m = 'i'
      # astk profile
      pr_astk = os.path.join(run['flasheur'], '%s.p%s' % (nomjob, njob))
      # if it doesn't exist the job is ended
      etat = "ENDED"
      if os.path.isfile(pr_astk):
         prof = ASTER_PROFIL(pr_astk, run)
         node = prof['noeud'][0]
   
   else:
      run.Mess(_('unexpected mode : %s') % mode, _('<F>_UNEXPECTED_VALUE'))

   # 2. query the process
   if node != '_':
      if node.split('.')[0] == run.GetHostName().split('.')[0]:
         jret, psout = run.Shell(run['ps_cpu'])
      else:
         jret, psout = run.Shell(run['ps_cpu'], mach=node)
      # ended ?
      if mode == "interactif" and psout.find('btc.%s' % njob)>-1:
         etat = "RUN"

   # 3.1. the job is ended
   if etat == "ENDED":
      fdiag = os.path.join(run['flasheur'], '%s.%s%s' % (nomjob, m, njob))
      if os.path.isfile(fdiag):
         diag = open(fdiag, 'r').read().split(os.linesep)[0]
         if diag == '':
            diag = '?'
      if diag == '?' :
         diag = _('<F>_SYSTEM')
         # try to find something in output
         fout = os.path.join(run['flasheur'], '%s.o%s' % (nomjob, njob))
         if os.path.isfile(fout):
            f = open(fout, 'r')
            for line in f:
               if line.find('--- DIAGNOSTIC JOB :')>-1:
                  diag = line.split()[4]
               elif line.find('Cputime limit exceeded')>-1:
                  diag = _('<F>_CPU_LIMIT_SYSTEM')
            f.close()
         # working directory
         wrk = _nodepara(run, node, 'rep_trav')
         # copy fort.6 to '.o'
         if not node in ('', '_'):
            ftcp = os.path.join(run['tmp_user'], 'as_actu.%s' % os.getpid())
            fort6 = os.path.join(wrk, '%s.%s.fort.6.%s' % (nomjob, njob, m))
            jret = run.Copy(ftcp, '%s:%s' % (node, fort6), niverr='SILENT')
            if os.path.isfile(ftcp):
               txt = [os.linesep*2]
               txt.append('='*48)
               txt.append('===== Pas de diagnostic, recopie du fort.6 =====')
               txt.append('='*48)
               txt.append(open(ftcp, 'r').read())
               txt.append('='*48)
               txt.append('='*48)
               txt.append(os.linesep*2)
               f = open(fout, 'a')
               f.write(os.linesep.join(txt))
               f.close()
   
   else:
   # 3.2. job is running
      if etat in ('RUN', 'SUSPENDED'):
         # working directory
         wrk = _nodepara(run, node, 'rep_trav')
         if not on_machref or mode == 'interactif':
            wrk = os.path.join(wrk, '%s.%s' % (mode, njob))
         else:
            wrk = os.path.join(wrk, str(jobid))
      if etat == 'RUN' and tcpu == '_':
         # tcpu may have been retrieved upper
         l_tcpu = []
         for line in psout.split(os.linesep):
            if re.search('\-num_job +%s' % njob, line) != None and \
               re.search('\-mode +%s' % mode, line) != None and \
               (not on_machref or line.find(run['rep_agla'])<0):
               l_tcpu.append(re.sub('\..*$', '', line.split()[0]).replace('-', ':'))
         if len(l_tcpu)>0:
            tcpu = _sumtps(l_tcpu)
   
   # 4. return the result
   return etat, diag, node, tcpu, wrk, queue
   
#-------------------------------------------------------------------------------
def Actu(run, *args):
   """Return state, diagnosis, node, cpu time and working directory of a job.
   """
   print "ETAT=%s DIAG=%s EXEC=%s TCPU=%s REP_TRAV=%s QUEUE=%s" % Func_actu(run, *args)

#-------------------------------------------------------------------------------
def Edit(run, *args):
   """Open output or error file of a job.
   """
   if len(args) != 5:
      run.parser.error(_("'--%s' takes exactly %d arguments (%d given)") % \
         (run.current_action, 5, len(args)))

   njob, nomjob, mode, typ, displ = args
   
   # get editor command line
   edit = run['editeur']
   
   # filename to edit
   fich = os.path.join(run['flasheur'], '%s.%s%s' % (nomjob, typ[0], njob))
   
   # open the file
   if os.path.isfile(fich):
      if re.search('display +@D', edit) != None:
         edit = edit.replace('@D', displ)
         kret, out = run.Shell('%s %s' % (edit, fich), bg=True)
      else:
         kret, out = run.Shell('%s -display %s %s' % (edit, displ, fich), bg=True)
   else:
      run.Sortie(4)

#-------------------------------------------------------------------------------
def Del(run, *args):
   """Kill a job and delete related files.
   """
   if len(args) < 3:
      run.parser.error(_("'--%s' takes at least %d arguments (%d given)") % \
         (run.current_action, 3, len(args)))
   elif len(args) > 4:
      run.parser.error(_("'--%s' takes at most %d arguments (%d given)") % \
         (run.current_action, 4, len(args)))
   
   # 0. arguments
   njob, nomjob, mode = args[:3]
   if len(args) > 3:
      node = args[3]
   else:
      node = ''
   
   use_batch_cmd = (mode == "batch" and run['signal'] == 'KILL')
      
   # 1. retrieve the job status
   etat, diag, node, tcpu, wrk, queue = Func_actu(run, njob, nomjob, mode)
   
   # 2. send the signal
   if etat in ('RUN', 'SUSPENDED', 'PEND'):
      if use_batch_cmd:
         if run['batch_nom'] in ('LSF', 'PBS', 'SunGE'):
            cmd = ''
            if run.get('batch_ini', '') != '':
               cmd = '. %s ; ' % run['batch_ini']
            cmd += '%s %s' % (run['batch_kil'], njob)
            iret, out = run.Shell(cmd)
            if iret != 0:
               run.Sortie(4)
         else:
            run.Mess(_('unknown batch scheduler'), _('<F>_ERROR'))
      else:
         numpr = ''
         # get process id
         if node != '_':
            if node.split('.')[0] == run.GetHostName().split('.')[0]:
               jret, psout = run.Shell(run['ps_pid'])
            else:
               jret, psout = run.Shell(run['ps_pid'], mach=node)
            exp = re.compile('^ *([0-9]+) +(.*)\-num_job +%s.*\-mode +%s' % (njob, mode),
                  re.MULTILINE)
            res = exp.findall(psout)
            for numj, cmd in res:
               if cmd.find('/bin/sh') < 0:
                  numpr = int(numj)
                  break
         if numpr != '':
            if node.split('.')[0] == run.GetHostName().split('.')[0]:
               os.kill(numpr, getattr(signal, 'SIG%s' % run['signal']))
            else:
               iret, psout = run.Shell('kill -%s %s' % (run['signal'], numpr), mach=node)
         else:
            run.Mess(_('<as_job.Del> process found :'))
            run.Mess(psout, 'SILENT')
            run.Mess(_('process not found'), _('<F>_ERROR'))
   
   # 3. delete files
   if run['signal'] == 'KILL':
      l_fich = glob.glob(os.path.join(run['flasheur'], '%s.?%s' % (nomjob, njob)))
      for f in l_fich:
         run.Delete(f)
   

#-------------------------------------------------------------------------------
def Tail(run, *args):
   """Output the last part of fort.6 file or filter lines matching a pattern.
   """
   if len(args) < 5:
      run.parser.error(_("'--%s' takes at least %d arguments (%d given)") % \
         (run.current_action, 5, len(args)))
   elif len(args) > 6:
      run.parser.error(_("'--%s' takes at most %d arguments (%d given)") % \
         (run.current_action, 6, len(args)))
   
   # 0. arguments
   njob, nomjob, mode, fdest, nbline = args[:5]
   if len(args) > 5:
      cmd = 'egrep -- \'%s\' ' % args[5]
   else:
      cmd = 'tail -%s ' % nbline
   
   # 1. retrieve the job status
   etat, diag, node, tcpu, wrk, queue = Func_actu(run, njob, nomjob, mode)

   # 2. fill output file
   s_out = ''
   ftmp = os.path.join(run['tmp_user'], 'as_tail.%s' % os.getpid())
   if mode == 'batch' and run['batch_nom'] == 'SunGE':
      s_out = _("Sorry I don't know how to ask Sun Grid Engine batch " \
                "scheduler the running node.")
   if etat == 'RUN':
      # 2.1.1. file to parse
      fich = os.path.join(wrk, 'fort.6')
      if node != '_':
         mach = node
         fich = '%s:%s' % (node, fich)
      else:
         mach = ''
      # 2.1.2. execute command
      if run.Exists(fich):
         jret, s_out = run.Shell(cmd + run.PathOnly(fich), mach=node)
      else:
         run.Mess(_('file is empty'))
   else:
      # 2.2. job isn't running
      print "JOB=%s JOBID=%s ETAT=%s" % (nomjob, njob, etat)
      run.Sortie(4)
   
   # 3. send output file
   if run.IsRemote(fdest):
      open(ftmp, 'w').write(s_out)
      jret = run.Copy(fdest, ftmp)
   else:
      fdest = run.PathOnly(fdest)
      open(fdest, 'w').write(s_out)

#-------------------------------------------------------------------------------
def Purge(run, *args):
   """Delete files which are NOT listed in args.
   """
   if len(args) < 1:
      run.parser.error(_("'--%s' takes at least %d arguments (%d given)") % \
         (run.current_action, 1, len(args)))
   
   # 0. arguments
   njob = len(args)
   l_job = []
   for j in args:
      try:
         l_job.append(int(j))
      except ValueError:
         run.parser.error(_('incorrect job number : %s') % j)

   # 1. action
   l_keep = []
   l_f = os.listdir(run['flasheur'])
   for f in l_f:
      delete = True
      for j in l_job:
         if re.search('\.[a-z]?%d$' % j, f) != None:
            delete = False
            break
      if delete:
         run.Delete(os.path.join(run['flasheur'], f), verbose=True)

