# @package      hubzero-submit-monitors
# @file         MonitorJobSite.py
# @copyright    Copyright (c) 2025-2025 The Regents of the University of California.
# @license      http://opensource.org/licenses/MIT MIT
#
# Copyright (c) 2025-2025 The Regents of the University of California.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# HUBzero is a registered trademark of The Regents of the University of California.
#
import sys
import os
import subprocess
import select
import signal
import time
import json
import traceback
import logging
from errno import EPIPE,ECONNRESET,EINTR

from hubzero.submit.LogMessage            import getLogPIDMessage as getLogMessage
from hubzero.submit.DaemonsInfo           import DaemonsInfo
from hubzero.submit.InfosInfo             import InfosInfo
from hubzero.submit.MonitorsInfo          import MonitorsInfo
from hubzero.submit.RemoteTunnelMonitor   import RemoteTunnelMonitor
from hubzero.submit.RemoteIdentityManager import RemoteIdentityManager
from hubzero.submit.UnboundConnection     import UnboundConnection

class MonitorJobSite:
   def __init__(self,
                daemonsDirectory,
                daemonsConfigurationFile,
                infosDirectory,
                infosConfigurationFile,
                identityUser,
                siteDesignator,
                activityUpdateInterval):
      self.logger                 = logging.getLogger(__name__)

      self.daemonsDirectory         = daemonsDirectory
      self.daemonsConfigurationFile = daemonsConfigurationFile
      self.infosDirectory           = infosDirectory
      self.infosConfigurationFile   = infosConfigurationFile
      self.activityUpdateInterval   = activityUpdateInterval

      configFilePath         = os.path.join(self.daemonsDirectory,self.daemonsConfigurationFile)
      daemonsInfo            = DaemonsInfo(configFilePath)
      self.tunnelListenURI   = daemonsInfo.getDaemonListenURI('tunnelMonitor','tcp')
      self.jobListenURI      = daemonsInfo.getDaemonListenURI('jobMonitor','tcp')
      self.identityListenURI = daemonsInfo.getDaemonListenURI('identitiesManager','tcp')

      self.identityUser   = identityUser
      self.siteDesignator = siteDesignator

      self.bufferSize     = 4096
      self.siteMonitorPid = 0

      self.infosInfo           = None
      self.monitorsInfo        = None
      self.remoteTunnelMonitor = None

      self.jobMonitorConnection = UnboundConnection(UnboundConnection.TLSREQUIREMENTNONE,[self.jobListenURI])

      self.exitSystemCallOnInterrupt = False

      signal.signal(signal.SIGINT,self.sigINT_handler)
      signal.signal(signal.SIGHUP,self.sigHUP_handler)
      signal.signal(signal.SIGQUIT,self.sigQUIT_handler)
      signal.signal(signal.SIGABRT,self.sigABRT_handler)
      signal.signal(signal.SIGTERM,self.sigTERM_handler)


   def terminate(self):
      if self.siteMonitorPid:
         self.logger.log(logging.INFO,getLogMessage("Send TERM to child ssh process"))
         try:
            os.kill(self.siteMonitorPid,signal.SIGTERM)
         except:
            pass


   def sigGEN_handler(self,
                      signalNumber,
                      frame):
      self.terminate()
      if self.exitSystemCallOnInterrupt:
         if sys.version_info > (3,):
            raise InterruptedError(EINTR,"System Call Interrupt")


   def sigINT_handler(self,
                      signalNumber,
                      frame):
      self.logger.log(logging.INFO,getLogMessage("Received SIGINT!"))
      self.sigGEN_handler(signalNumber,frame)


   def sigHUP_handler(self,
                      signalNumber,
                      frame):
      self.logger.log(logging.INFO,getLogMessage("Received SIGHUP!"))
      self.sigGEN_handler(signalNumber,frame)


   def sigQUIT_handler(self,
                       signalNumber,
                       frame):
      self.logger.log(logging.INFO,getLogMessage("Received SIGQUIT!"))
      self.sigGEN_handler(signalNumber,frame)


   def sigABRT_handler(self,
                       signalNumber,
                       frame):
      self.logger.log(logging.INFO,getLogMessage("Received SIGABRT!"))
      self.sigGEN_handler(signalNumber,frame)


   def sigTERM_handler(self,
                       signalNumber,
                       frame):
      self.logger.log(logging.INFO,getLogMessage("Received SIGTERM!"))
      self.sigGEN_handler(signalNumber,frame)


   def writeToSiteMonitor(self,
                          childin,
                          jsonMessage):
      try:
         childin.write((jsonMessage + '\n').encode('utf-8'))
         childin.flush()
      except Exception as e:
         if isinstance(e.args,tuple):
            if   e.args[0] == EPIPE:
               self.logger.log(logging.ERROR,getLogMessage("writeToSiteMonitor: Broken pipe"))
            elif e.args[0] == ECONNRESET:
               self.logger.log(logging.ERROR,getLogMessage("writeToSiteMonitor: Connection reset by peer"))
            else:
               self.logger.log(logging.ERROR,getLogMessage("Unexpected error in writeToSiteMonitor(%s)" % (jsonMessage[0:50])))
               self.logger.log(logging.ERROR,getLogMessage(traceback.format_exc()))
         else:
            self.logger.log(logging.ERROR,getLogMessage("Unexpected error in writeToSiteMonitor(%s)" % (jsonMessage[0:50])))
            self.logger.log(logging.ERROR,getLogMessage(traceback.format_exc()))
         self.terminate()


   def monitorSite(self):
      exitStatus = 0
      if self.jobMonitorConnection.isConnected():
         configFilePath = os.path.join(self.infosDirectory,self.infosConfigurationFile)
         self.infosInfo             = InfosInfo(configFilePath)
         self.monitorsInfo          = MonitorsInfo(self.infosInfo.getInfoPath('monitors'))
         self.remoteTunnelMonitor   = RemoteTunnelMonitor(self.tunnelListenURI)
         self.remoteIdentityManager = RemoteIdentityManager(self.identityListenURI)

         monitorInfo = self.monitorsInfo.getMonitorInfo(self.siteDesignator)
         sshCommand       = []
         tunnelDesignator = ""
         if monitorInfo:
            identityPaths = self.remoteIdentityManager.queryUserIdentities(monitorInfo['identityManager'],
                                                                           self.identityUser)
            try:
               monitorName = monitorInfo['monitorName']
               if monitorInfo['venueMechanism'] == 'ssh':
                  venue                = monitorInfo['venue']
                  venuePort            = monitorInfo['venuePort']
                  sshOptions           = monitorInfo['sshOptions']
                  tunnelDesignator     = monitorInfo['tunnelDesignator']
                  user                 = monitorInfo['remoteUser']
                  remoteMonitorCommand = monitorInfo['remoteMonitorCommand']
                  sshIdentity          = identityPaths['communitySSH']

                  sshCommand.append("ssh")
                  sshCommand.append("-T")
                  sshCommand.append("-x")
                  sshCommand.append("-a")
                  if sshOptions:
                     sshCommand += sshOptions.split()
                  sshCommand.append("-i")
                  sshCommand.append(sshIdentity)
                  if tunnelDesignator == "":
                     sshCommand.append("-p")
                     sshCommand.append("%d" % (venuePort))
                     sshCommand.append("%s@%s" % (user,venue))
                     sshCommand.append(remoteMonitorCommand)
                     sshCommand.append(monitorName)
                  else:
                     tunnelAddress,tunnelPort = self.remoteTunnelMonitor.getTunnelAddressPort(tunnelDesignator)
                     sshCommand.append("-p")
                     sshCommand.append(tunnelPort)
                     sshCommand.append("%s@%s" % (user,tunnelAddress))
                     sshCommand.append(remoteMonitorCommand)
                     sshCommand.append(monitorName)
            except:
               self.logger.log(logging.ERROR,getLogMessage("Build sshCommand failed for %s" % (self.siteDesignator)))
               self.logger.log(logging.ERROR,getLogMessage(traceback.format_exc()))

         if sshCommand:
            self.logger.log(logging.INFO,getLogMessage("****************************************"))
            self.logger.log(logging.INFO,getLogMessage("* %s job monitor started *" % (self.siteDesignator)))
            self.logger.log(logging.INFO,getLogMessage("****************************************"))

            if tunnelDesignator != "":
               self.remoteTunnelMonitor.incrementTunnelUse(tunnelDesignator)

            child = subprocess.Popen(sshCommand,bufsize=1,
                                     stdin=subprocess.PIPE,
                                     stdout=subprocess.PIPE,
                                     stderr=subprocess.PIPE,
                                     close_fds=True)
            self.siteMonitorPid = child.pid
            childin    = child.stdin
            childinFd  = childin.fileno()
            childout   = child.stdout
            childoutFd = childout.fileno()
            childerr   = child.stderr
            childerrFd = childerr.fileno()
            outEOF = False
            errEOF = False
            errData = []

            jsonObject = {'messageType':'startJobSite',
                          'siteDesignator':self.siteDesignator,
                          'pid':os.getpid()}
            self.jobMonitorConnection.postJsonMessage(jsonObject)

            lastTimeFakeNewRemoteJobId = 0
            newRemoteJobIds = []
            while True:
               readers = self.jobMonitorConnection.getInputObject()
               if not outEOF:
                  readers.append(childoutFd)
               if not errEOF:
                  readers.append(childerrFd)
               writers = self.jobMonitorConnection.getOutputObject()
               if not outEOF or not errEOF:
                  now = time.time()
                  if len(newRemoteJobIds) > 0:
                     writers.append(childinFd)
                     lastTimeFakeNewRemoteJobId = now
                  else:
                     if now-lastTimeFakeNewRemoteJobId > self.activityUpdateInterval/5.:
                        writers.append(childinFd)
                        lastTimeFakeNewRemoteJobId = now

               if outEOF and errEOF and not writers:
                  break

               try:
                  self.exitSystemCallOnInterrupt = True
                  readyReaders,readyWriters,readyExceptions = select.select(readers,writers,[],self.activityUpdateInterval)
               except select.error as e:
                  if e.args[0] == EINTR:
                     readyReaders = []
                     readyWriters = []
                  else:
                     self.terminate()
               finally:
                  self.exitSystemCallOnInterrupt = False

               for readyReader in readyReaders:
                  if   readyReader == childoutFd:
#                    self.logger.log(logging.DEBUG,getLogMessage("LOADING MESSAGE"))
                     jsonMessage = childout.readline().decode('utf-8')
                     if jsonMessage == '':
                        outEOF = True
                     else:
#                       self.logger.log(logging.INFO,getLogMessage("%d character message received from %s" % \
#                                                                     (len(jsonMessage),self.siteDesignator)))
                        while jsonMessage != "":
                           try:
                              jsonObject = json.loads(jsonMessage)
                           except ValueError:
                              self.logger.log(logging.ERROR,getLogMessage("JSON object %s could not be decoded" % (jsonMessage)))
                           else:
                              if jsonObject['messageType'] != 'pipeFlusher':
                                 self.logger.log(logging.DEBUG,getLogMessage("request = %s, message length = %d" % \
                                                                      (jsonObject['messageType'],len(jsonMessage))))
                                 requestStart = time.time()
                              if   jsonObject['messageType'] == 'siteUpdate':
                                 messages = jsonObject['jobStates']
                                 messageLength = jsonObject['nJobStates']

                                 self.logger.log(logging.INFO,getLogMessage("%d/%d job status messages received from %s" % \
                                                                         (len(messages),messageLength,self.siteDesignator)))
                                 self.jobMonitorConnection.postJsonMessage(jsonObject)
                              elif jsonObject['messageType'] == 'pipeFlusher':
                                 self.jobMonitorConnection.postJsonMessage(jsonObject)
                                 break
                              else:
                                 self.logger.log(logging.ERROR,getLogMessage("Invalid message from site %s" % \
                                                                                        (self.siteDesignator)))
                              if jsonObject['messageType'] != 'pipeFlusher':
                                 requestTime = time.time()-requestStart
                                 self.logger.log(logging.DEBUG,getLogMessage("request = %s, processing time = %f" % \
                                                                            (jsonObject['messageType'],requestTime)))
#                          self.logger.log(logging.DEBUG,getLogMessage("LOADING NEXT MESSAGE"))
                           jsonMessage = childout.readline().decode('utf-8')
                  elif readyReader == childerrFd:
                     errChunk = os.read(childerrFd,self.bufferSize).decode('utf-8')
                     if errChunk == '':
                        errEOF = True
                     errData.append(errChunk)
                  else:
                     self.jobMonitorConnection.receiveMessage()

               message = self.jobMonitorConnection.pullMessage(0)
               while message:
                  args = message.split()
                  if args[0] == 'json':
                     jsonMessageLength = int(args[1])
                     jsonMessage = self.jobMonitorConnection.pullMessage(jsonMessageLength)
                     if len(jsonMessage) > 0:
                        try:
                           jsonObject = json.loads(jsonMessage)
                        except ValueError:
                           self.logger.log(logging.ERROR,getLogMessage("JSON object %s could not be decoded" % (jsonMessage)))
                        else:
                           if jsonObject['messageType'] == 'newJobId':
                              newRemoteJobIds.append(jsonObject)
                           else:
                              self.logger.log(logging.ERROR,getLogMessage("Discarded message type: %s" % \
                                                                             (jsonObject['messageType'])))
                     else:
                        self.jobMonitorConnection.pushMessage(message + '\n')
                        break
                  else:
                     self.logger.log(logging.ERROR,getLogMessage("Discarded message: %s" % (message)))

                  message = self.jobMonitorConnection.pullMessage(0)

               for readyWriter in readyWriters:
                  if readyWriter == childinFd:
                     for jsonObject in newRemoteJobIds:
                        try:
                           jsonMessage = json.dumps(jsonObject)
                        except TypeError:
                           self.logger.log(logging.ERROR,getLogMessage("JSON object %s could not be encoded" % (jsonObject)))
                        else:
                           self.writeToSiteMonitor(childin,jsonMessage)
                     newRemoteJobIds = []

                     jsonObject = {'messageType':'pipeFlusher'}
                     try:
                        jsonMessage = json.dumps(jsonObject)
                     except:
                        pass
                     else:
#                       self.logger.log(logging.DEBUG,getLogMessage("Site pipe flushed"))
                        self.writeToSiteMonitor(childin,jsonMessage)
                     time.sleep(5)
                  else:
                     self.jobMonitorConnection.sendMessage()

               if outEOF and errEOF and self.siteMonitorPid:
                  pid,exitStatus = os.waitpid(self.siteMonitorPid,0)
                  self.siteMonitorPid = 0
                  if exitStatus != 0:
                     if   os.WIFSIGNALED(exitStatus):
                        self.logger.log(logging.INFO,getLogMessage("%s failed w/ signal %d" % (sshCommand,os.WTERMSIG(exitStatus))))
                     else:
                        if os.WIFEXITED(exitStatus):
                           exitStatus = os.WEXITSTATUS(exitStatus)
                        self.logger.log(logging.INFO,getLogMessage("%s failed w/ exit code %d" % (sshCommand,exitStatus)))
                     self.logger.log(logging.INFO,getLogMessage("%s" % ("".join(errData))))

                  if tunnelDesignator != "":
                     self.remoteTunnelMonitor.decrementTunnelUse(tunnelDesignator)

                  jsonObject = {'messageType':'stopJobSite',
                                'siteDesignator':self.siteDesignator,
                                'pid':os.getpid(),
                                'exitStatus':exitStatus}
                  self.jobMonitorConnection.postJsonMessage(jsonObject)

            self.logger.log(logging.INFO,getLogMessage("****************************************"))
            self.logger.log(logging.INFO,getLogMessage("* %s job monitor stopped *" % (self.siteDesignator)))
            self.logger.log(logging.INFO,getLogMessage("****************************************"))
         else:
            exitStatus = 1
      else:
         self.logger.log(logging.ERROR,getLogMessage("Connection to central job monitor failed"))
         exitStatus = 1

      return(exitStatus)


