#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Implementation of ``DataChannel``s to communicate across the data plane."""

# pytype: skip-file

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

import abc
import collections
import logging
import queue
import sys
import threading
import time
from builtins import object
from builtins import range
from typing import TYPE_CHECKING
from typing import Callable
from typing import DefaultDict
from typing import Dict
from typing import Iterable
from typing import Iterator
from typing import List
from typing import Optional
from typing import Set
from typing import Tuple
from typing import Type
from typing import Union

import grpc
from future.utils import raise_
from future.utils import with_metaclass

from apache_beam.coders import coder_impl
from apache_beam.portability.api import beam_fn_api_pb2
from apache_beam.portability.api import beam_fn_api_pb2_grpc
from apache_beam.runners.worker.channel_factory import GRPCChannelFactory
from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor

if TYPE_CHECKING:
  # TODO(BEAM-9372): move this out of the TYPE_CHECKING scope when we drop
  #  support for python < 3.5.3
  from types import TracebackType
  ExcInfo = Tuple[Type[BaseException], BaseException, TracebackType]
  OptExcInfo = Union[ExcInfo, Tuple[None, None, None]]
  # TODO: move this out of the TYPE_CHECKING scope when we drop support for
  #  python < 3.6
  from typing import Collection  # pylint: disable=ungrouped-imports
  import apache_beam.coders.slow_stream
  OutputStream = apache_beam.coders.slow_stream.OutputStream
  DataOrTimers = \
    Union[beam_fn_api_pb2.Elements.Data, beam_fn_api_pb2.Elements.Timers]
else:
  OutputStream = type(coder_impl.create_OutputStream())

# This module is experimental. No backwards-compatibility guarantees.

_LOGGER = logging.getLogger(__name__)

_DEFAULT_SIZE_FLUSH_THRESHOLD = 10 << 20  # 10MB
_DEFAULT_TIME_FLUSH_THRESHOLD_MS = 0  # disable time-based flush by default


class ClosableOutputStream(OutputStream):
  """A Outputstream for use with CoderImpls that has a close() method."""
  def __init__(self, close_callback=None):
    super(ClosableOutputStream, self).__init__()
    self._close_callback = close_callback

  def close(self):
    # type: () -> None
    if self._close_callback:
      self._close_callback(self.get())

  @staticmethod
  def create(close_callback, flush_callback, data_buffer_time_limit_ms):
    # type: (...) -> SizeBasedBufferingClosableOutputStream
    if data_buffer_time_limit_ms > 0:
      return TimeBasedBufferingClosableOutputStream(
          close_callback,
          flush_callback=flush_callback,
          time_flush_threshold_ms=data_buffer_time_limit_ms)
    else:
      return SizeBasedBufferingClosableOutputStream(
          close_callback, flush_callback=flush_callback)


class SizeBasedBufferingClosableOutputStream(ClosableOutputStream):
  """A size-based buffering OutputStream."""

  def __init__(self,
               close_callback=None,  # type: Optional[Callable[[bytes], None]]
               flush_callback=None,  # type: Optional[Callable[[bytes], None]]
               size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD):
    super(SizeBasedBufferingClosableOutputStream, self).__init__(close_callback)
    self._flush_callback = flush_callback
    self._size_flush_threshold = size_flush_threshold

  # This must be called explicitly to avoid flushing partial elements.
  def maybe_flush(self):
    if self.size() > self._size_flush_threshold:
      self.flush()

  def flush(self):
    if self._flush_callback:
      self._flush_callback(self.get())
      self._clear()


class TimeBasedBufferingClosableOutputStream(
    SizeBasedBufferingClosableOutputStream):
  """A buffering OutputStream with both time-based and size-based."""
  def __init__(
      self,
      close_callback=None,
      flush_callback=None,
      size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD,
      time_flush_threshold_ms=_DEFAULT_TIME_FLUSH_THRESHOLD_MS):
    super(TimeBasedBufferingClosableOutputStream,
          self).__init__(close_callback, flush_callback, size_flush_threshold)
    assert time_flush_threshold_ms > 0
    self._time_flush_threshold_ms = time_flush_threshold_ms
    self._flush_lock = threading.Lock()
    self._schedule_lock = threading.Lock()
    self._closed = False
    self._schedule_periodic_flush()

  def flush(self):
    with self._flush_lock:
      super(TimeBasedBufferingClosableOutputStream, self).flush()

  def close(self):
    with self._schedule_lock:
      self._closed = True
      if self._periodic_flusher:
        self._periodic_flusher.cancel()
        self._periodic_flusher = None
    super(TimeBasedBufferingClosableOutputStream, self).close()

  def _schedule_periodic_flush(self):
    def _flush():
      with self._schedule_lock:
        if not self._closed:
          self.flush()

    self._periodic_flusher = PeriodicThread(
        self._time_flush_threshold_ms / 1000.0, _flush)
    self._periodic_flusher.daemon = True
    self._periodic_flusher.start()


class PeriodicThread(threading.Thread):
  """Call a function periodically with the specified number of seconds"""
  def __init__(self, interval, function, args=None, kwargs=None):
    threading.Thread.__init__(self)
    self._interval = interval
    self._function = function
    self._args = args if args is not None else []
    self._kwargs = kwargs if kwargs is not None else {}
    self._finished = threading.Event()

  def run(self):
    next_call = time.time() + self._interval
    while not self._finished.wait(next_call - time.time()):
      next_call = next_call + self._interval
      self._function(*self._args, **self._kwargs)

  def cancel(self):
    """Stop the thread if it hasn't finished yet."""
    self._finished.set()


class DataChannel(with_metaclass(abc.ABCMeta, object)):  # type: ignore[misc]
  """Represents a channel for reading and writing data over the data plane.

  Read data and timer from this channel with the input_elements method::

    for elements_data in data_channel.input_elements(
        instruction_id, transform_ids, timers):
      [process elements_data]

  Write data to this channel using the output_stream method::

    out1 = data_channel.output_stream(instruction_id, transform_id)
    out1.write(...)
    out1.close()

  Write timer to this channel using the output_timer_stream method::

    out1 = data_channel.output_timer_stream(instruction_id,
                                            transform_id,
                                            timer_family_id)
    out1.write(...)
    out1.close()

  When all data/timer for all instructions is written, close the channel::

    data_channel.close()
  """
  @abc.abstractmethod
  def input_elements(self,
                     instruction_id,  # type: str
                     expected_inputs,  # type: Collection[str]
                     abort_callback=None  # type: Optional[Callable[[], bool]]
                    ):
    # type: (...) -> Iterator[DataOrTimers]

    """Returns an iterable of all Element.Data and Element.Timers bundles for
    instruction_id.

    This iterable terminates only once the full set of data has been recieved
    for each of the expected transforms. It may block waiting for more data.

    Args:
        instruction_id: which instruction the results must belong to
        expected_transforms: which transforms to wait on for completion
        abort_callback: a callback to invoke if blocking returning whether
            to abort before consuming all the data
    """
    raise NotImplementedError(type(self))

  @abc.abstractmethod
  def output_stream(
      self,
      instruction_id,  # type: str
      transform_id  # type: str
  ):
    # type: (...) -> ClosableOutputStream

    """Returns an output stream writing elements to transform_id.

    Args:
        instruction_id: which instruction this stream belongs to
        transform_id: the transform_id of the returned stream
    """
    raise NotImplementedError(type(self))

  @abc.abstractmethod
  def output_timer_stream(self,
                              instruction_id,  # type: str
                              transform_id,  # type: str
                              timer_family_id  # type: str
                          ):
    # type: (...) -> ClosableOutputStream

    """Returns an output stream written timers to transform_id.

    Args:
        instruction_id: which instruction this stream belongs to
        transform_id: the transform_id of the returned stream
        timer_family_id: the timer family of the written timer
    """
    raise NotImplementedError(type(self))

  @abc.abstractmethod
  def close(self):
    # type: () -> None

    """Closes this channel, indicating that all data has been written.

    Data can continue to be read.

    If this channel is shared by many instructions, should only be called on
    worker shutdown.
    """
    raise NotImplementedError(type(self))


class InMemoryDataChannel(DataChannel):
  """An in-memory implementation of a DataChannel.

  This channel is two-sided.  What is written to one side is read by the other.
  The inverse() method returns the other side of a instance.
  """
  def __init__(self, inverse=None, data_buffer_time_limit_ms=0):
    # type: (Optional[InMemoryDataChannel], Optional[int]) -> None
    self._inputs = []  # type: List[DataOrTimers]
    self._data_buffer_time_limit_ms = data_buffer_time_limit_ms
    self._inverse = inverse or InMemoryDataChannel(
        self, data_buffer_time_limit_ms=data_buffer_time_limit_ms)

  def inverse(self):
    # type: () -> InMemoryDataChannel
    return self._inverse

  def input_elements(self,
      instruction_id,  # type: str
      unused_expected_inputs=None,   # type: Collection[str]
      abort_callback=None  # type: Optional[Callable[[], bool]]
                     ):
    other_inputs = []
    for element in self._inputs:
      if element.instruction_id == instruction_id:
        if isinstance(element, beam_fn_api_pb2.Elements.Timers):
          if not element.is_last:
            yield element
        if isinstance(element, beam_fn_api_pb2.Elements.Data):
          if element.data or element.is_last:
            yield element
      else:
        other_inputs.append(element)
    self._inputs = other_inputs

  def output_timer_stream(self, instruction_id, transform_id, timer_family_id):
    def add_to_inverse_output(timer):
      if timer:
        self._inverse._inputs.append(
            beam_fn_api_pb2.Elements.Timers(
                instruction_id=instruction_id,
                transform_id=transform_id,
                timer_family_id=timer_family_id,
                timers=timer,
                is_last=False))

    def close_stream(timer):
      add_to_inverse_output(timer)
      self._inverse._inputs.append(
          beam_fn_api_pb2.Elements.Timers(
              instruction_id=instruction_id,
              transform_id=transform_id,
              timer_family_id='',
              is_last=True))

    return ClosableOutputStream.create(
        add_to_inverse_output, close_stream, self._data_buffer_time_limit_ms)

  def output_stream(self, instruction_id, transform_id):
    # type: (str, str) -> ClosableOutputStream
    def add_to_inverse_output(data):
      self._inverse._inputs.append(  # pylint: disable=protected-access
          beam_fn_api_pb2.Elements.Data(
              instruction_id=instruction_id,
              transform_id=transform_id,
              data=data))

    return ClosableOutputStream.create(
        add_to_inverse_output,
        add_to_inverse_output,
        self._data_buffer_time_limit_ms)

  def close(self):
    pass


class _GrpcDataChannel(DataChannel):
  """Base class for implementing a BeamFnData-based DataChannel."""

  _WRITES_FINISHED = object()

  def __init__(self, data_buffer_time_limit_ms=0):
    # type: (Optional[int]) -> None
    self._data_buffer_time_limit_ms = data_buffer_time_limit_ms
    self._to_send = queue.Queue()  # type: queue.Queue[DataOrTimers]
    self._received = collections.defaultdict(
        lambda: queue.Queue(maxsize=5)
    )  # type: DefaultDict[str, queue.Queue[DataOrTimers]]

    self._receive_lock = threading.Lock()
    self._reads_finished = threading.Event()
    self._closed = False
    self._exc_info = None  # type: Optional[OptExcInfo]

  def close(self):
    self._to_send.put(self._WRITES_FINISHED)
    self._closed = True

  def wait(self, timeout=None):
    self._reads_finished.wait(timeout)

  def _receiving_queue(self, instruction_id):
    # type: (str) -> queue.Queue[DataOrTimers]
    with self._receive_lock:
      return self._received[instruction_id]

  def _clean_receiving_queue(self, instruction_id):
    # type: (str) -> None
    with self._receive_lock:
      self._received.pop(instruction_id)

  def input_elements(self,
      instruction_id,  # type: str
      expected_inputs,   # type: Collection[Union[str, Tuple[str, str]]]
      abort_callback=None  # type: Optional[Callable[[], bool]]
                     ):

    # type: (...) -> Iterator[DataOrTimers]

    """
    Generator to retrieve elements for an instruction_id
    input_elements should be called only once for an instruction_id

    Args:
      instruction_id(str): instruction_id for which data is read
      expected_inputs(collection): expected inputs, include both data and timer.
    """
    received = self._receiving_queue(instruction_id)
    done_inputs = set()  # type: Set[Union[str, Tuple[str, str]]]
    abort_callback = abort_callback or (lambda: False)
    try:
      while len(done_inputs) < len(expected_inputs):
        try:
          element = received.get(timeout=1)
        except queue.Empty:
          if self._closed:
            raise RuntimeError('Channel closed prematurely.')
          if abort_callback():
            return
          if self._exc_info:
            t, v, tb = self._exc_info
            raise_(t, v, tb)
        else:
          if isinstance(element, beam_fn_api_pb2.Elements.Timers):
            if element.is_last:
              done_inputs.add((element.transform_id, element.timer_family_id))
            else:
              yield element
          elif isinstance(element, beam_fn_api_pb2.Elements.Data):
            if element.is_last:
              done_inputs.add(element.transform_id)
            else:
              assert element.transform_id not in done_inputs
              yield element
          else:
            raise ValueError('Unexpected input element type %s' % type(element))
    finally:
      # Instruction_ids are not reusable so Clean queue once we are done with
      #  an instruction_id
      self._clean_receiving_queue(instruction_id)

  def output_stream(self, instruction_id, transform_id):
    # type: (str, str) -> ClosableOutputStream
    def add_to_send_queue(data):
      # type: (bytes) -> None
      if data:
        self._to_send.put(
            beam_fn_api_pb2.Elements.Data(
                instruction_id=instruction_id,
                transform_id=transform_id,
                data=data))

    def close_callback(data):
      # type: (bytes) -> None
      add_to_send_queue(data)
      # End of stream marker.
      self._to_send.put(
          beam_fn_api_pb2.Elements.Data(
              instruction_id=instruction_id,
              transform_id=transform_id,
              is_last=True))

    return ClosableOutputStream.create(
        close_callback, add_to_send_queue, self._data_buffer_time_limit_ms)

  def output_timer_stream(self, instruction_id, transform_id, timer_family_id):
    def add_to_send_queue(timer):
      if timer:
        self._to_send.put(
            beam_fn_api_pb2.Elements.Timers(
                instruction_id=instruction_id,
                transform_id=transform_id,
                timer_family_id=timer_family_id,
                timers=timer,
                is_last=False))

    def close_callback(timer):
      add_to_send_queue(timer)
      self._to_send.put(
          beam_fn_api_pb2.Elements.Timers(
              instruction_id=instruction_id,
              transform_id=transform_id,
              timer_family_id=timer_family_id,
              is_last=True))

    return ClosableOutputStream.create(
        close_callback, add_to_send_queue, self._data_buffer_time_limit_ms)

  def _write_outputs(self):
    # type: () -> Iterator[beam_fn_api_pb2.Elements]
    stream_done = False
    while not stream_done:
      streams = [self._to_send.get()]
      try:
        # Coalesce up to 100 other items.
        for _ in range(100):
          streams.append(self._to_send.get_nowait())
      except queue.Empty:
        pass
      if streams[-1] is self._WRITES_FINISHED:
        stream_done = True
        streams.pop()
      if streams:
        data_stream = []
        timer_stream = []
        for stream in streams:
          if isinstance(stream, beam_fn_api_pb2.Elements.Timers):
            timer_stream.append(stream)
          elif isinstance(stream, beam_fn_api_pb2.Elements.Data):
            data_stream.append(stream)
          else:
            raise ValueError('Unexpected output element type %s' % type(stream))
        yield beam_fn_api_pb2.Elements(data=data_stream, timers=timer_stream)

  def _read_inputs(self, elements_iterator):
    # type: (Iterable[beam_fn_api_pb2.Elements]) -> None
    try:
      for elements in elements_iterator:
        for timer in elements.timers:
          self._receiving_queue(timer.instruction_id).put(timer)
        for data in elements.data:
          self._receiving_queue(data.instruction_id).put(data)
    except:  # pylint: disable=bare-except
      if not self._closed:
        _LOGGER.exception('Failed to read inputs in the data plane.')
        self._exc_info = sys.exc_info()
        raise
    finally:
      self._closed = True
      self._reads_finished.set()

  def set_inputs(self, elements_iterator):
    # type: (Iterable[beam_fn_api_pb2.Elements]) -> None
    reader = threading.Thread(
        target=lambda: self._read_inputs(elements_iterator),
        name='read_grpc_client_inputs')
    reader.daemon = True
    reader.start()


class GrpcClientDataChannel(_GrpcDataChannel):
  """A DataChannel wrapping the client side of a BeamFnData connection."""

  def __init__(self,
               data_stub,  # type: beam_fn_api_pb2_grpc.BeamFnDataStub
               data_buffer_time_limit_ms=0  # type: Optional[int]
               ):
    # type: (...) -> None
    super(GrpcClientDataChannel, self).__init__(data_buffer_time_limit_ms)
    self.set_inputs(data_stub.Data(self._write_outputs()))


class BeamFnDataServicer(beam_fn_api_pb2_grpc.BeamFnDataServicer):
  """Implementation of BeamFnDataServicer for any number of clients"""
  def __init__(
      self,
      data_buffer_time_limit_ms=0  # type: Optional[int]
  ):
    self._lock = threading.Lock()
    self._connections_by_worker_id = collections.defaultdict(
        lambda: _GrpcDataChannel(data_buffer_time_limit_ms)
    )  # type: DefaultDict[str, _GrpcDataChannel]

  def get_conn_by_worker_id(self, worker_id):
    # type: (str) -> _GrpcDataChannel
    with self._lock:
      return self._connections_by_worker_id[worker_id]

  def Data(self,
           elements_iterator,  # type: Iterable[beam_fn_api_pb2.Elements]
           context
          ):
    # type: (...) -> Iterator[beam_fn_api_pb2.Elements]
    worker_id = dict(context.invocation_metadata())['worker_id']
    data_conn = self.get_conn_by_worker_id(worker_id)
    data_conn.set_inputs(elements_iterator)
    for elements in data_conn._write_outputs():
      yield elements


class DataChannelFactory(with_metaclass(abc.ABCMeta,
                                        object)):  # type: ignore[misc]
  """An abstract factory for creating ``DataChannel``."""
  @abc.abstractmethod
  def create_data_channel(self, remote_grpc_port):
    # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel

    """Returns a ``DataChannel`` from the given RemoteGrpcPort."""
    raise NotImplementedError(type(self))

  @abc.abstractmethod
  def close(self):
    # type: () -> None

    """Close all channels that this factory owns."""
    raise NotImplementedError(type(self))


class GrpcClientDataChannelFactory(DataChannelFactory):
  """A factory for ``GrpcClientDataChannel``.

  Caches the created channels by ``data descriptor url``.
  """

  def __init__(self,
               credentials=None,
               worker_id=None,  # type: Optional[str]
               data_buffer_time_limit_ms=0  # type: Optional[int]
               ):
    # type: (...) -> None
    self._data_channel_cache = {}  # type: Dict[str, GrpcClientDataChannel]
    self._lock = threading.Lock()
    self._credentials = None
    self._worker_id = worker_id
    self._data_buffer_time_limit_ms = data_buffer_time_limit_ms
    if credentials is not None:
      _LOGGER.info('Using secure channel creds.')
      self._credentials = credentials

  def create_data_channel_from_url(self, url):
    if not url:
      return None
    if url not in self._data_channel_cache:
      with self._lock:
        if url not in self._data_channel_cache:
          _LOGGER.info('Creating client data channel for %s', url)
          # Options to have no limits (-1) on the size of the messages
          # received or sent over the data plane. The actual buffer size
          # is controlled in a layer above.
          channel_options = [("grpc.max_receive_message_length", -1),
                             ("grpc.max_send_message_length", -1)]
          grpc_channel = None
          if self._credentials is None:
            grpc_channel = GRPCChannelFactory.insecure_channel(
                url, options=channel_options)
          else:
            grpc_channel = GRPCChannelFactory.secure_channel(
                url, self._credentials, options=channel_options)
          # Add workerId to the grpc channel
          grpc_channel = grpc.intercept_channel(
              grpc_channel, WorkerIdInterceptor(self._worker_id))
          self._data_channel_cache[url] = GrpcClientDataChannel(
              beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel),
              self._data_buffer_time_limit_ms)

    return self._data_channel_cache[url]

  def create_data_channel(self, remote_grpc_port):
    # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel
    url = remote_grpc_port.api_service_descriptor.url
    return self.create_data_channel_from_url(url)

  def close(self):
    # type: () -> None
    _LOGGER.info('Closing all cached grpc data channels.')
    for _, channel in self._data_channel_cache.items():
      channel.close()
    self._data_channel_cache.clear()


class InMemoryDataChannelFactory(DataChannelFactory):
  """A singleton factory for ``InMemoryDataChannel``."""
  def __init__(self, in_memory_data_channel):
    # type: (GrpcClientDataChannel) -> None
    self._in_memory_data_channel = in_memory_data_channel

  def create_data_channel(self, unused_remote_grpc_port):
    # type: (beam_fn_api_pb2.RemoteGrpcPort) -> GrpcClientDataChannel
    return self._in_memory_data_channel

  def create_data_channel_from_url(self, url):
    return self._in_memory_data_channel

  def close(self):
    # type: () -> None
    pass
