# Copyright 2017 The TensorFlow Authors. 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.
# ==============================================================================
"""Wrapper around debugger plugin to conditionally enable it."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys

import six
from werkzeug import wrappers

from tensorboard.backend import http_util
from tensorboard.plugins import base_plugin
from tensorboard.plugins.debugger import constants
from tensorboard.util import tb_logging

logger = tb_logging.get_logger()


class InactiveDebuggerPlugin(base_plugin.TBPlugin):
  """A placeholder debugger plugin used when no grpc port is specified."""

  plugin_name = constants.DEBUGGER_PLUGIN_NAME

  def __init__(self):
    pass

  def is_active(self):
    return False

  def frontend_metadata(self):
    return base_plugin.FrontendMetadata(element_name='tf-debugger-dashboard')

  def get_plugin_apps(self):
    return {
        '/debugger_grpc_host_port': self._serve_debugger_grpc_host_port,
    }

  @wrappers.Request.application
  def _serve_debugger_grpc_host_port(self, request):
    # Respond with a -1 port number to indicate the debugger plugin is
    # inactive.
    return http_util.Respond(
        request,
        {'host': None, 'port': -1},
        'application/json')



class DebuggerPluginLoader(base_plugin.TBLoader):
  """DebuggerPlugin factory factory.

  This class determines which debugger plugin to load, based on custom
  flags. It also checks for the `grpcio` PyPi dependency.
  """

  def define_flags(self, parser):
    """Adds DebuggerPlugin CLI flags to parser."""
    group = parser.add_argument_group('debugger plugin')
    group.add_argument(
        '--debugger_data_server_grpc_port',
        metavar='PORT',
        type=int,
        default=-1,
        help='''\
The port at which the non-interactive debugger data server should
receive debugging data via gRPC from one or more debugger-enabled
TensorFlow runtimes. No debugger plugin or debugger data server will be
started if this flag is not provided. This flag differs from the
`--debugger_port` flag in that it starts a non-interactive mode. It is
for use with the "health pills" feature of the Graph Dashboard. This
flag is mutually exclusive with `--debugger_port`.\
''')
    group.add_argument(
        '--debugger_port',
        metavar='PORT',
        type=int,
        default=-1,
        help='''\
The port at which the interactive debugger data server (to be started by
the debugger plugin) should receive debugging data via gRPC from one or
more debugger-enabled TensorFlow runtimes. No debugger plugin or
debugger data server will be started if this flag is not provided. This
flag differs from the `--debugger_data_server_grpc_port` flag in that it
starts an interactive mode that allows user to pause at selected nodes
inside a TensorFlow Graph or between Session.runs. It is for use with
the interactive Debugger Dashboard. This flag is mutually exclusive with
`--debugger_data_server_grpc_port`.\
''')

  def fix_flags(self, flags):
    """Fixes Debugger related flags.

    Raises:
      ValueError: If both the `debugger_data_server_grpc_port` and
        `debugger_port` flags are specified as >= 0.
    """
    # Check that not both grpc port flags are specified.
    if flags.debugger_data_server_grpc_port > 0 and flags.debugger_port > 0:
      raise base_plugin.FlagsError(
          '--debugger_data_server_grpc_port and --debugger_port are mutually '
          'exclusive. Do not use both of them at the same time.')

  def load(self, context):
    """Returns the debugger plugin, if possible.

    Args:
      context: The TBContext flags including `add_arguments`.

    Returns:
      A DebuggerPlugin instance or None if it couldn't be loaded.
    """
    flags = context.flags
    if flags.debugger_data_server_grpc_port > 0 or flags.debugger_port > 0:
      # Verify that the required Python packages are installed.
      try:
        # pylint: disable=g-import-not-at-top,unused-import
        import tensorflow
      except ImportError:
        raise ImportError(
            'To use the debugger plugin, you need to have TensorFlow installed:\n'
            '  pip install tensorflow')

    if flags.debugger_data_server_grpc_port > 0:
      # pylint: disable=line-too-long,g-import-not-at-top
      from tensorboard.plugins.debugger import debugger_plugin as debugger_plugin_lib
      # pylint: enable=line-too-long,g-import-not-at-top

      # debugger_data_server_grpc opens the non-interactive Debugger Plugin,
      # which appears as health pills in the Graph Plugin.
      noninteractive_plugin = debugger_plugin_lib.DebuggerPlugin(context)
      logger.info('Starting Non-interactive Debugger Plugin at gRPC port %d',
                   flags.debugger_data_server_grpc_port)
      noninteractive_plugin.listen(flags.debugger_data_server_grpc_port)
      return noninteractive_plugin
    elif flags.debugger_port > 0:
      # pylint: disable=line-too-long,g-import-not-at-top
      from tensorboard.plugins.debugger import interactive_debugger_plugin as interactive_debugger_plugin_lib
      # pylint: enable=line-too-long,g-import-not-at-top
      interactive_plugin = (
          interactive_debugger_plugin_lib.InteractiveDebuggerPlugin(context))
      logger.info('Starting Interactive Debugger Plugin at gRPC port %d',
                  flags.debugger_data_server_grpc_port)
      interactive_plugin.listen(flags.debugger_port)
      return interactive_plugin
    else:
      # If neither the debugger_data_server_grpc_port flag or the grpc_port
      # flag is specified, we instantiate a dummy plugin as a placeholder for
      # the frontend. The dummy plugin will display a message indicating that
      # the plugin is not active. It'll also display a command snippet to
      # illustrate how to activate the interactive Debugger Plugin.
      return InactiveDebuggerPlugin()
