#!/usr/bin/env python

"""
Launch Jupyter Notebook within a Docker notebook image and 
automatically open up the URL within the default web browser.
"""

# Author: Xiangmin Jiao <xiangmin.jiao@stonybrook.edu>

from __future__ import print_function # Only Python 2.x

import argparse
import random
import string
import sys

# Process command-line arguments
parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument('-u', "--user", nargs='?', help='''username & groupname used by the image. The default is numgeom.''', default="numgeom")

parser.add_argument('-i', '--image', nargs='?', help='''The Docker image to use The default is pdepack-notebook.''', default="pdepack-notebook")

parser.add_argument('notebook', nargs='?', help='''The notebook to open.''', default="")

args=parser.parse_args()
image = args.image
user = args.user
notebook = args.notebook

container = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(12))

def random_ports(port, n):
    """Generate a list of n random ports near the given port.

    The first 5 ports will be sequential, and the remaining n-5 will be
    randomly selected in the range [port-2*n, port+2*n].
    """
    import random

    for i in range(min(5, n)):
        yield port + i
    for i in range(n-5):
        yield max(1, port + random.randint(-2*n, 2*n))

def find_free_port(port, retries):
   "Find a free port"
   import socket

   sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

   for port in random_ports(port, retries+1):
      result = sock.connect_ex(('127.0.0.1', port))
      if result:
         return port

   print("Error: Could not find a free port")
   sys.exit(-1)

if __name__ == "__main__":
   import os
   import subprocess
   import webbrowser

   pwd = os.getcwd()

   # Find an unused port
   port = str(find_free_port(8888, 50))

   # Start the docker image and pipe the stderr
   p = subprocess.Popen(["docker", "run", "--name", container,
                         "-w", "/home/" + user + "/shared", 
                         "-v", pwd + ":/home/" + user + "/shared", 
                         "-p", port + ":" + port, 
                         "quay.io/" + user +"/" + image, 
                         "jupyter-notebook --no-browser --ip=0.0.0.0 --port " + port], 
                        stderr=subprocess.PIPE, universal_newlines=True)
   
   # Wait for user to press Ctrl-C twice
   try:
      opened = False
      # Monitor the stderr to extract the URL
      for stderr_line in iter(p.stderr.readline, ""):
         ind = stderr_line.find("http://0.0.0.0:")
         sys.stderr.write(stderr_line)

         if not opened and ind >= 0:
            # Open browser if found URL
            if not notebook:
                webbrowser.open("http://localhost:" + stderr_line[ind+15:])
            else:
                webbrowser.open("http://localhost:" + port + "/notebooks/" + notebook + 
                                stderr_line[stderr_line.find("?token="):])
                
            opened = True

      p.stderr.close()
      p.wait()
   except KeyboardInterrupt:
       # Terminate docker if user pressed Ctrl-C
       while True:
           try:
               tosave = input('''Press Ctrl-C again to stop. To save your changes to the
Docker image before exiting, type "commit":''')

               if tosave.find("commit") >= 0:
                   # Commit the image
                   subprocess.call(["docker", "commit", container, "quay.io/" + user +"/" + image]);
                   # Terminate the process and remove the docker container
                   p.terminate()
                   subprocess.call(["docker", "rm", "-f", container]);
                   break
               else:
                   continue
           except KeyboardInterrupt:
               # Terminate the process and remove the docker container
               p.terminate()
               subprocess.call(["docker", "rm", "-f", container]);
               break
   except OSError:
       pass
