# Copyright 2019 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.
# ==============================================================================
"""Dumping op callbacks: Enables dump-based features in tfdbg v2."""

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

import atexit
import re
import socket
import threading
import uuid

from six.moves import xrange  # pylint: disable=redefined-builtin

from tensorflow.core.protobuf import debug_event_pb2
from tensorflow.core.protobuf import graph_debug_info_pb2
from tensorflow.python.debug.lib import debug_events_writer
from tensorflow.python.debug.lib import op_callbacks_common
from tensorflow.python.debug.lib import source_utils
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import op_callbacks
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_debug_ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.util import compat
from tensorflow.python.util import tf_stack
from tensorflow.python.util.tf_export import tf_export

_state = threading.local()
DEFAULT_TENSOR_DEBUG_MODE = "NO_TENSOR"


@ops.RegisterGradient("DebugIdentityV2")
def _debug_identity_v2_grad(op, dy):
  """Gradient function for the DebugIdentityV2 op."""
  del op  # Unused
  return dy


def _get_id():
  """Get a short unique ID."""
  return str(uuid.uuid4())


class _DumpingCallback(object):
  """An object holding the states surrouding the dumping callback."""

  def __init__(self,
               dump_root,
               tensor_debug_mode,
               circular_buffer_size,
               op_regex,