import argparse
import time
import logging
import threading
import signal
import socket
import subprocess

FORMAT="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
logging.basicConfig(level = logging.DEBUG, format = FORMAT)

def kill_gentally(proc):
  proc.terminate()
  time.sleep(1)
  if proc.poll() is None:
    logging.info("SIGTERM not working, try again with SIGKILL")
    proc.kill()

class Handle:

  def __init__(self, args):
    self.state = "INIT"
    self.args = args
    self.proc = None
    self.killFlag = False
    self.exitFlag = False

  def restart(self):
    if self.state == 'STARTED':
      self.killFlag = True

  def exit(self):
    self.exitFlag = True

  def loop(self):
    self._start()
    while True:
      if self.exitFlag:
        self._kill()
        break
      elif self.proc.poll() is not None: # died
        self.state = "INITED"
        retcode = self.proc.poll()
        logging.info("[died] process died with code: %d", retcode)
        self._start()
      elif self.killFlag:
        self._kill()
        self._start()
      time.sleep(1)
      pass

  def _kill(self):
    if self.state != 'STARTED':
      logging.warn("_kill is called in state: %s", self.state)
      return
    self.state = "KILLING"
    if self.proc.poll() is None:
      kill_gentally(self.proc)
    retcode = self.proc.wait()
    logging.info("[killed] process died with code: %d", retcode)
    self.state = "INIT"

  def _start(self):
    self.killFlag = False
    self.state = "STARTING"
    logging.info("starting ssh...")
    self.proc = subprocess.Popen(['ssh'] + self.args)
    self.state = "STARTED"

def start_ssh_loop_in_thread(sshargs):
  handle = Handle(sshargs)
  thread = threading.Thread(target=handle.loop)
  #thread.setDaemon(True)
  thread.start()
  handle.join = thread.join
  return handle

def check_monitor(port):
  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  try:
    sock.connect(('127.0.0.1', port))
    sock.close()
    return True
  except:
    return False

def main(opts):
  handle = start_ssh_loop_in_thread(opts.sshargs)
  flags = {'exit': False}
  def exit_gracefully(signum, frame):
    handle.exit()
    flags['exit'] = True

  signal.signal(signal.SIGINT, exit_gracefully)
  signal.signal(signal.SIGTERM, exit_gracefully)
  print opts.sshargs, opts.monitor
  while flags['exit'] == False:
    if not check_monitor(opts.monitor):
      logging.error("check monitor failed")
      handle.restart()
    time.sleep(10)
  handle.join()
  logging.info("fully exited!")

if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument("monitor", type=int)
  #parser.add_argument('-', action='store_true')
  #parser.add_argument('--daemon', action = 'store_true')
  #parser.add_argument('--pidfile', default=os.path.join(_my_dir_, _my_file_ + ".pid"))
  #parser.add_argument('--logfile', default = os.path.join(_my_dir_, _my_file_ + ".log"))
  parser.add_argument('sshargs', nargs=argparse.REMAINDER)
  main(parser.parse_args())