# Copyright 2016 Google Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Methods for implementing the `datalab connect` command."""

from __future__ import absolute_import

import os
import threading
import time
import webbrowser

try:
    from urllib.request import urlopen
except ImportError:
    from urllib2 import urlopen

from . import utils


description = """`{0} {1}` creates a persistent connection to a
Datalab instance running in a Google Compute Engine VM.

This is a thin wrapper around the *ssh(1)* command that takes care
of authentication and the translation of the instance name into an
IP address.

This command ensures that the user's public SSH key is present
in the project's metadata. If the user does not have a public
SSH key, one is generated using *ssh-keygen(1)* (if the --quiet
flag is given, the generated key will have an empty passphrase).

Once the connection is established, this command opens a browser
window pointing at Datalab. To disable the browser opening, add
the --no-launch-browser flag.

This command will attempt to re-establish the connection if it
gets dropped. However, that connection will only exist while
this command is running.
"""


examples = """
To connect to 'example-instance' in zone 'us-central1-a', run:

    $ {0} {1} example-instance --zone us-central1-a"""


wrong_user_message_template = (
    'The specified Datalab instance was created for {0}, but you '
    'are attempting to connect to it as {1}.'
    '\n\n'
    'Datalab instances are single-user environments, and trying '
    'to share one is not supported.'
    '\n\n'
    'To override this behavior, re-run the command with the '
    '--no-user-checking flag.')


web_preview_message_template = (
    'Click on the *Web Preview* (square button at top-right), '
    'select *Change port > Port {}*, and start using Datalab.')


connection_closed_message_template = """
Connection closed.

To re-connect to your datalab instance, run the following command:

    datalab connect{1}{0}
"""


# The list of web browsers that we don't want to automatically open.
#
# This is a subset of the canonical list of python browser types
# defined here: https://docs.python.org/2/library/webbrowser.html
unsupported_browsers = [
    'BackgroundBrowser', 'Elinks', 'GenericBrowser', 'Grail'
]


# Status values from describe_instance that we care about.
_STATUS_RUNNING = 'RUNNING'


def flags(parser):
    """Add command line flags for the `connect` subcommand.

    Args:
      parser: The argparse parser to which to add the flags.
    """
    parser.add_argument(
        'instance',
        metavar='NAME',
        help='name of the instance to which to connect')
    parser.add_argument(
        '--no-user-checking',
        dest='no_user_checking',
        action='store_true',
        default=False,
        help='do not check if the current user matches the Datalab instance')

    connection_flags(parser)
    return


def connection_flags(parser):
    """Add flags common to every connection-establishing subcommand.

    Args:
      parser: The argparse parser to which to add the flags.
    """
    parser.add_argument(
        '--port',
        dest='port',
        type=int,
        default=8081,
        help='local port on which Datalab is accessible')
    parser.add_argument(
        '--max-reconnects',
        dest='max_reconnects',
        type=int,
        default=-1,
        help=(
            'maximum number of times to reconnect.'
            '\n\n'
            'A negative value means no limit.'))
    parser.add_argument(
        '--connection-health-timeout-seconds',
        dest='connection_health_timeout_seconds',
        type=int,
        default=120,
        help=(
            'maximum amount of time, in seconds, to wait for a connection '
            'to become healthy before giving up and retrying.'))
    parser.add_argument(
        '--ssh-log-level',
        dest='ssh_log_level',
        choices=['quiet', 'fatal', 'error', 'info', 'verbose',
                 'debug', 'debug1', 'debug2', 'debug3'],
        default='error',
        help=(
            'the log level for the SSH command.'
            '\n\n'
            'This may be useful for debugging issues with an SSH connection.'
            '\n\n'
            'The default log level is "error".'
            '\n\n'
            'This has no effect when run on Windows.'))
    parser.add_argument(
        '--no-launch-browser',
        dest='no_launch_browser',
        action='store_true',
        default=False,
        help='do not open a browser connected to Datalab')
    parser.add_argument(
        '--beta-internal-ip',
        dest='internal_ip',
        action='store_true',
        default=False,
        help=('connect to the internal IP address of the instance.'
              '\n\n'
              'Note that this is a beta feature and unsupported.'))

    return


def connect(args, gcloud_compute, email, in_cloud_shell):
    """Create a persistent connection to a Datalab instance.

    Args:
      args: The Namespace object constructed by argparse
      gcloud_compute: A function that can be called to invoke `gcloud compute`
      email: The user's email address
      in_cloud_shell: Whether or not the command is being run in the
        Google Cloud Shell
    """
    instance = args.instance
    connect_msg = ('Connecting to {0}.\n'
                   'This will create an SSH tunnel '
                   'and may prompt you to create an rsa key pair. '
                   'To manage these keys, see https://cloud.google.com/'
                   'compute/docs/instances/adding-removing-ssh-keys')
    print(connect_msg.format(instance))

    datalab_port = args.port
    datalab_address = 'http://localhost:{0}/'.format(str(datalab_port))

    def create_tunnel():
        """Create an SSH tunnel to the Datalab instance.

        This method blocks for as long as the connection is open.

        Raises:
          KeyboardInterrupt: When the end user kills the connection
        """
        if utils.print_debug_messages(args):
            print('Connecting to {0} via SSH').format(instance)

        cmd = ['ssh']