#!/home/patrickd/swarms/swarmsenv/bin/python3.12
#
# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
# Copyright (c) 1997-2016 California Institute of Technology.
# Copyright (c) 2016-2024 The Uncertainty Quantification Foundation.
# License: 3-clause BSD.  The full license text is available at:
#  - https://github.com/uqfoundation/pathos/blob/master/LICENSE
"""
connect to the specified machine and start a 'server', 'tunnel', or both

Notes:
    Usage: pathos_connect [hostname] [server] [remoteport] [profile]
        [hostname] - name of the host to connect to
        [server] - name of RPC server (assumes is installed on host) or 'tunnel'
        [remoteport] - remote port to use for communication or 'tunnel'
        [profile] -- name of shell profile to source on remote environment

Examples::

    $ pathos_connect computer.college.edu ppserver tunnel
    Usage: pathos_connect [hostname] [server] [remoteport] [profile]
        [hostname] - name of the host to connect to
        [server] - name of RPC server (assumes is installed on host) or 'tunnel'
        [remoteport] - remote port to use for communication or 'tunnel'
        [profile] -- name of shell profile to source on remote environment
        defaults are: "localhost" "tunnel" "" ""
    executing {ssh -N -L 22921:computer.college.edu:15058}'

    Server running at port=15058 with pid=4110
    Connected to localhost at port=22921
    Press <Enter> to kill server
"""
## tunnel: pathos_connect college.edu tunnel
## server: pathos_connect college.edu ppserver 12345 .profile
## both:   pathos_connect college.edu ppserver tunnel .profile

from pathos.core import *
from pathos.hosts import get_profile, register_profiles


if __name__ == '__main__':

##### CONFIGURATION & INPUT ########################
  # set the default remote host
  rhost = 'localhost'
 #rhost = 'foobar.internet.org'
 #rhost = 'computer.college.edu'

  # set any 'special' profiles (those which don't use default_profie)
  profiles = {}
 #profiles = {'foobar.internet.org':'.profile',
 #            'computer.college.edu':'.cshrc'}

  # set the default port
  rport = ''
  _rport = '98909'

  # set the default server command
  server = 'tunnel'
 #server = 'ppserver'  #XXX: "ppserver -p %s" % rport
 #server = 'classic_server'  #XXX: "classic_server -p %s" % rport
 #server = 'registry_server'  #XXX: "registry_server -p %s" % rport

  print("""Usage: pathos_connect [hostname] [remoteport] [server] [profile]
    Usage: pathos_connect [hostname] [server] [remoteport] [profile]
    [hostname] - name of the host to connect to
    [server] - name of RPC server (assumes is installed on host) or 'tunnel'
    [remoteport] - remote port to use for communication or 'tunnel'
    [profile] -- name of shell profile to source on remote environment
    defaults are: "%s" "%s" "%s" "%s".""" % (rhost, server, rport, ''))

  # get remote hostname from user
  import sys
  if '--help' in sys.argv:
    sys.exit(0)
  try:
    myinp = sys.argv[1]
  except: myinp = None
  if myinp:
    rhost = myinp #XXX: should test rhost validity here... (how ?)
  else: pass # use default
  del myinp

  # get server to run from user
  try:
    myinp = sys.argv[2]
  except: myinp = None
  if myinp:
    server = myinp #XXX: should test validity here... (filename)
  else: pass # use default
  del myinp

  # set the default 'port'
  if server == 'tunnel':
    tunnel = True
    server = None
  else:
    tunnel = False
  rport = rport if tunnel else _rport

  # get remote port to run server on from user
  try:
    myinp = sys.argv[3]
  except: myinp = None
  if myinp:
    if tunnel: # tunnel doesn't take more inputs
      msg = "port '%s' not valid for 'tunnel'" % myinp
      raise ValueError(msg)
    rport = myinp #XXX: should test validity here... (filename)
  else: pass # use default
  del myinp

  # is it a tunneled server?
  tunnel = True if (tunnel or rport == 'tunnel') else False 
  rport = '' if rport == 'tunnel' else rport

  # get remote profile (this should go away soon)
  try:
    myinp = sys.argv[4]
  except: myinp = None
  if myinp:
    rprof = myinp #XXX: should test validity here... (filename)
    profiles = {rhost:rprof}
  else: pass # use default
  del myinp

  # my remote environment (should be auto-detected)
  register_profiles(profiles)
  profile = get_profile(rhost)

##### CONFIGURATION & INPUT ########################
## tunnel: pathos_connect foo.college.edu tunnel
## server: pathos_connect foo.college.edu ppserver 12345 .profile
## both:   pathos_connect foo.college.edu ppserver tunnel .profile

  if tunnel:
    # establish ssh tunnel
    tunnel = connect(rhost)
    lport = tunnel._lport
    rport = tunnel._rport
    print('executing {ssh -N -L %d:%s:%d}' % (lport, rhost, rport))
  else:
    lport = ''

  if server:
    # run server
    rserver = serve(server, rhost, rport, profile=profile)
    response = rserver.response()
    if response:
      if tunnel: tunnel.disconnect()
      print(response)
      raise OSError('Failure to start server')

    # get server pid  #FIXME: launcher.pid is not pid(server)
    target = '[P,p]ython[^#]*'+server # filter w/ regex for python-based server
    try:
      pid = getpid(target, rhost)
    except OSError:
      print("Cleanup on host may be required...")
      if tunnel: tunnel.disconnect()
      raise

    # test server
    # XXX: add a simple one-liner...
    print("\nServer running at port=%s with pid=%s" % (rport, pid))
    if tunnel: print("Connected to localhost at port=%s" % (lport))
    print('Press <Enter> to kill server')
  else:
    print('Press <Enter> to disconnect')
  sys.stdin.readline()

  if server:
    # stop server
    print(kill(pid,rhost))
#   del rserver  #XXX: delete should run self.kill (?)

  if tunnel:
    # disconnect tunnel
    tunnel.disconnect()
    # FIXME: just kills 'ssh', not the tunnel
    # get local pid: ps u | grep "ssh -N -L%s:%s$s" % (lport,rhost,rport)
    # kill -15 int(tunnelpid)

# EOF