#
# 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.
#
"""SDK harness for executing Python Fns via the Fn API."""

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

import abc
import collections
import contextlib
import logging
import queue
import sys
import threading
import time
import traceback
from builtins import object
from builtins import range
from concurrent import futures

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

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 import bundle_processor
from apache_beam.runners.worker import data_plane
from apache_beam.runners.worker.channel_factory import GRPCChannelFactory
from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor


class SdkHarness(object):
  REQUEST_METHOD_PREFIX = '_request_'
  SCHEDULING_DELAY_THRESHOLD_SEC = 5*60  # 5 Minutes

  def __init__(
      self, control_address, worker_count, credentials=None, worker_id=None,
      profiler_factory=None):
    self._alive = True
    self._worker_count = worker_count
    self._worker_index = 0
    self._worker_id = worker_id
    if credentials is None:
      logging.info('Creating insecure control channel for %s.', control_address)
      self._control_channel = GRPCChannelFactory.insecure_channel(
          control_address)
    else:
      logging.info('Creating secure control channel for %s.', control_address)
      self._control_channel = GRPCChannelFactory.secure_channel(
          control_address, credentials)
    grpc.channel_ready_future(self._control_channel).result(timeout=60)
    logging.info('Control channel established.')

    self._control_channel = grpc.intercept_channel(
        self._control_channel, WorkerIdInterceptor(self._worker_id))
    self._data_channel_factory = data_plane.GrpcClientDataChannelFactory(
        credentials)
    self._state_handler_factory = GrpcStateHandlerFactory(credentials)
    self._profiler_factory = profiler_factory
    self._fns = {}
    # BundleProcessor cache across all workers.
    self._bundle_processor_cache = BundleProcessorCache(
        state_handler_factory=self._state_handler_factory,
        data_channel_factory=self._data_channel_factory,
        fns=self._fns)
    # workers for process/finalize bundle.
    self.workers = queue.Queue()
    # one worker for progress/split request.
    self.progress_worker = SdkWorker(self._bundle_processor_cache,
                                     profiler_factory=self._profiler_factory)
    # one thread is enough for getting the progress report.
    # Assumption:
    # Progress report generation should not do IO or wait on other resources.
    #  Without wait, having multiple threads will not improve performance and
    #  will only add complexity.
    self._progress_thread_pool = futures.ThreadPoolExecutor(max_workers=1)
    # finalize and process share one thread pool.
    self._process_thread_pool = futures.ThreadPoolExecutor(
        max_workers=self._worker_count)
    self._responses = queue.Queue()
    self._process_bundle_queue = queue.Queue()
    self._unscheduled_process_bundle = {}
    logging.info('Initializing SDKHarness with %s workers.', self._worker_count)

  def run(self):
    control_stub = beam_fn_api_pb2_grpc.BeamFnControlStub(self._control_channel)
    no_more_work = object()

    # Create process workers
    for _ in range(self._worker_count):
      # SdkHarness manage function registration and share self._fns with all
      # the workers. This is needed because function registration (register)
      # and exceution(process_bundle) are send over different request and we
      # do not really know which woker is going to process bundle
      # for a function till we get process_bundle request. Moreover
      # same function is reused by different process bundle calls and
      # potentially get executed by different worker. Hence we need a
      # centralized function list shared among all the workers.
      self.workers.put(
          SdkWorker(self._bundle_processor_cache,
                    profiler_factory=self._profiler_factory))

    def get_responses():
      while True:
        response = self._responses.get()
        if response is no_more_work:
          return
        yield response

    self._alive = True
    monitoring_thread = threading.Thread(target=self._monitor_process_bundle)
    monitoring_thread.daemon = True
    monitoring_thread.start()

    try:
      for work_request in control_stub.Control(get_responses()):
        logging.debug('Got work %s', work_request.instruction_id)
        request_type = work_request.WhichOneof('request')
        # Name spacing the request method with 'request_'. The called method
        # will be like self.request_register(request)
        getattr(self, SdkHarness.REQUEST_METHOD_PREFIX + request_type)(
            work_request)
    finally:
      self._alive = False

    logging.info('No more requests from control plane')
    logging.info('SDK Harness waiting for in-flight requests to complete')
    # Wait until existing requests are processed.
    self._progress_thread_pool.shutdown()
    self._process_thread_pool.shutdown()
    # get_responses may be blocked on responses.get(), but we need to return
    # control to its caller.
    self._responses.put(no_more_work)
    # Stop all the workers and clean all the associated resources
    self._data_channel_factory.close()
    self._state_handler_factory.close()
    logging.info('Done consuming work.')

  def _execute(self, task, request):
    try:
      response = task()
    except Exception:  # pylint: disable=broad-except
      traceback_string = traceback.format_exc()
      print(traceback_string, file=sys.stderr)
      logging.error(
          'Error processing instruction %s. Original traceback is\n%s\n',
          request.instruction_id, traceback_string)
      response = beam_fn_api_pb2.InstructionResponse(
          instruction_id=request.instruction_id, error=traceback_string)
    self._responses.put(response)

  def _request_register(self, request):

    def task():
      for process_bundle_descriptor in getattr(
          request, request.WhichOneof('request')).process_bundle_descriptor:
        self._fns[process_bundle_descriptor.id] = process_bundle_descriptor

      return beam_fn_api_pb2.InstructionResponse(
          instruction_id=request.instruction_id,
          register=beam_fn_api_pb2.RegisterResponse())

    self._execute(task, request)

  def _request_process_bundle(self, request):

    def task():
      # Take the free worker. Wait till a worker is free.
      worker = self.workers.get()
      # Get the first work item in the queue
      work = self._process_bundle_queue.get()
      self._unscheduled_process_bundle.pop(work.instruction_id, None)
      try:
        self._execute(lambda: worker.do_instruction(work), work)
      finally:
        # Put the worker back in the free worker pool
        self.workers.put(worker)
    # Create a task for each process_bundle request and schedule it
    self._process_bundle_queue.put(request)
    self._unscheduled_process_bundle[request.instruction_id] = time.time()
    self._process_thread_pool.submit(task)
    logging.debug(
        "Currently using %s threads." % len(self._process_thread_pool._threads))

  def _request_process_bundle_split(self, request):
    self._request_process_bundle_action(request)

  def _request_process_bundle_progress(self, request):
    self._request_process_bundle_action(request)

  def _request_process_bundle_action(self, request):

    def task():
      instruction_reference = getattr(
          request, request.WhichOneof('request')).instruction_reference
      # only process progress/split request when a bundle is in processing.
      if (instruction_reference in
          self._bundle_processor_cache.active_bundle_processors):
        self._execute(
            lambda: self.progress_worker.do_instruction(request), request)
      else:
        self._execute(lambda: beam_fn_api_pb2.InstructionResponse(
            instruction_id=request.instruction_id, error=(
                'Process bundle request not yet scheduled for instruction {}' if
                instruction_reference in self._unscheduled_process_bundle else
                'Unknown process bundle instruction {}').format(
                    instruction_reference)), request)

    self._progress_thread_pool.submit(task)

  def _request_finalize_bundle(self, request):

    def task():
      # Get one available worker.
      worker = self.workers.get()
      try:
        self._execute(
            lambda: worker.do_instruction(request), request)
      finally:
        # Put the worker back in the free worker pool.
        self.workers.put(worker)

    self._process_thread_pool.submit(task)

  def _monitor_process_bundle(self):
    """
    Monitor the unscheduled bundles and log if a bundle is not scheduled for
    more than SCHEDULING_DELAY_THRESHOLD_SEC.
    """
    while self._alive:
      time.sleep(SdkHarness.SCHEDULING_DELAY_THRESHOLD_SEC)
      # Check for bundles to be scheduled.
      if self._unscheduled_process_bundle:
        current_time = time.time()
        for instruction_id in self._unscheduled_process_bundle:
          request_time = None
          try:
            request_time = self._unscheduled_process_bundle[instruction_id]
          except KeyError:
            pass
          if request_time:
            scheduling_delay = current_time - request_time
            if scheduling_delay > SdkHarness.SCHEDULING_DELAY_THRESHOLD_SEC:
              logging.warn('Unable to schedule instruction %s for %s',
                           instruction_id, scheduling_delay)


class BundleProcessorCache(object):
  """A cache for ``BundleProcessor``s.

  ``BundleProcessor`` objects are cached by the id of their
  ``beam_fn_api_pb2.ProcessBundleDescriptor``.

  Attributes:
    fns (dict): A dictionary that maps bundle descriptor IDs to instances of
      ``beam_fn_api_pb2.ProcessBundleDescriptor``.
    state_handler_factory (``StateHandlerFactory``): Used to create state
      handlers to be used by a ``bundle_processor.BundleProcessor`` during
      processing.
    data_channel_factory (``data_plane.DataChannelFactory``)
    active_bundle_processors (dict): A dictionary, indexed by instruction IDs,
      containing ``bundle_processor.BundleProcessor`` objects that are currently
      active processing the corresponding instruction.
    cached_bundle_processors (dict): A dictionary, indexed by bundle processor
      id, of cached ``bundle_processor.BundleProcessor`` that are not currently
      performing processing.
  """

  def __init__(self, state_handler_factory, data_channel_factory, fns):
    self.fns = fns
    self.state_handler_factory = state_handler_factory
    self.data_channel_factory = data_channel_factory
    self.active_bundle_processors = {}
    self.cached_bundle_processors = collections.defaultdict(list)

  def register(self, bundle_descriptor):
    """Register a ``beam_fn_api_pb2.ProcessBundleDescriptor`` by its id."""
    self.fns[bundle_descriptor.id] = bundle_descriptor

  def get(self, instruction_id, bundle_descriptor_id):
    try:
      # pop() is threadsafe
      processor = self.cached_bundle_processors[bundle_descriptor_id].pop()
    except IndexError:
      processor = bundle_processor.BundleProcessor(
          self.fns[bundle_descriptor_id],
          self.state_handler_factory.create_state_handler(
              self.fns[bundle_descriptor_id].state_api_service_descriptor),
          self.data_channel_factory)
    self.active_bundle_processors[
        instruction_id] = bundle_descriptor_id, processor
    return processor

  def lookup(self, instruction_id):
    return self.active_bundle_processors.get(instruction_id, (None, None))[-1]

  def discard(self, instruction_id):
    self.active_bundle_processors[instruction_id][1].shutdown()
    del self.active_bundle_processors[instruction_id]

  def release(self, instruction_id):
    descriptor_id, processor = self.active_bundle_processors.pop(instruction_id)
    processor.reset()
    self.cached_bundle_processors[descriptor_id].append(processor)

  def shutdown(self):
    for instruction_id in self.active_bundle_processors:
      self.active_bundle_processors[instruction_id][1].shutdown()
      del self.active_bundle_processors[instruction_id]
    for cached_bundle_processors in self.cached_bundle_processors.values():
      while len(cached_bundle_processors) > 0:
        cached_bundle_processors.pop().shutdown()


class SdkWorker(object):

  def __init__(self, bundle_processor_cache, profiler_factory=None):
    self.bundle_processor_cache = bundle_processor_cache
    self.profiler_factory = profiler_factory

  def do_instruction(self, request):
    request_type = request.WhichOneof('request')
    if request_type:
      # E.g. if register is set, this will call self.register(request.register))
      return getattr(self, request_type)(getattr(request, request_type),
                                         request.instruction_id)
    else:
      raise NotImplementedError

  def register(self, request, instruction_id):
    """Registers a set of ``beam_fn_api_pb2.ProcessBundleDescriptor``s.

    This set of ``beam_fn_api_pb2.ProcessBundleDescriptor`` come as part of a
    ``beam_fn_api_pb2.RegisterRequest``, which the runner sends to the SDK
    worker before starting processing to register stages.
    """

    for process_bundle_descriptor in request.process_bundle_descriptor:
      self.bundle_processor_cache.register(process_bundle_descriptor)
    return beam_fn_api_pb2.InstructionResponse(
        instruction_id=instruction_id,
        register=beam_fn_api_pb2.RegisterResponse())

  def process_bundle(self, request, instruction_id):
    bundle_processor = self.bundle_processor_cache.get(
        instruction_id, request.process_bundle_descriptor_reference)
    try:
      with bundle_processor.state_handler.process_instruction_id(
          instruction_id):
        with self.maybe_profile(instruction_id):
          delayed_applications, requests_finalization = (
              bundle_processor.process_bundle(instruction_id))
          response = beam_fn_api_pb2.InstructionResponse(
              instruction_id=instruction_id,
              process_bundle=beam_fn_api_pb2.ProcessBundleResponse(
                  residual_roots=delayed_applications,
                  metrics=bundle_processor.metrics(),
                  monitoring_infos=bundle_processor.monitoring_infos(),
                  requires_finalization=requests_finalization))
      # Don't release here if finalize is needed.
      if not requests_finalization:
        self.bundle_processor_cache.release(instruction_id)
      return response
    except:  # pylint: disable=broad-except
      # Don't re-use bundle processors on failure.
      self.bundle_processor_cache.discard(instruction_id)
      raise

  def process_bundle_split(self, request, instruction_id):
    processor = self.bundle_processor_cache.lookup(
        request.instruction_reference)
    if processor:
      return beam_fn_api_pb2.InstructionResponse(
          instruction_id=instruction_id,
          process_bundle_split=processor.try_split(request))
    else:
      return beam_fn_api_pb2.InstructionResponse(
          instruction_id=instruction_id,
          error='Instruction not running: %s' % instruction_id)

  def process_bundle_progress(self, request, instruction_id):
    # It is an error to get progress for a not-in-flight bundle.
    processor = self.bundle_processor_cache.lookup(
        request.instruction_reference)
    return beam_fn_api_pb2.InstructionResponse(
        instruction_id=instruction_id,
        process_bundle_progress=beam_fn_api_pb2.ProcessBundleProgressResponse(
            metrics=processor.metrics() if processor else None,
            monitoring_infos=processor.monitoring_infos() if processor else []))

  def finalize_bundle(self, request, instruction_id):
    processor = self.bundle_processor_cache.lookup(
        request.instruction_reference)
    if processor:
      try:
        finalize_response = processor.finalize_bundle()
        self.bundle_processor_cache.release(request.instruction_reference)
        return beam_fn_api_pb2.InstructionResponse(
            instruction_id=instruction_id,
            finalize_bundle=finalize_response)
      except:
        self.bundle_processor_cache.discard(request.instruction_reference)
        raise
    else:
      return beam_fn_api_pb2.InstructionResponse(
          instruction_id=instruction_id,
          error='Instruction not running: %s' % instruction_id)

  def stop(self):
    self.bundle_processor_cache.shutdown()

  @contextlib.contextmanager
  def maybe_profile(self, instruction_id):
    if self.profiler_factory:
      profiler = self.profiler_factory(instruction_id)
      if profiler:
        with profiler:
          yield
      else:
        yield
    else:
      yield


class StateHandlerFactory(with_metaclass(abc.ABCMeta, object)):
  """An abstract factory for creating ``DataChannel``."""

  @abc.abstractmethod
  def create_state_handler(self, api_service_descriptor):
    """Returns a ``StateHandler`` from the given ApiServiceDescriptor."""
    raise NotImplementedError(type(self))

  @abc.abstractmethod
  def close(self):
    """Close all channels that this factory owns."""
    raise NotImplementedError(type(self))


class GrpcStateHandlerFactory(StateHandlerFactory):
  """A factory for ``GrpcStateHandler``.

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

  def __init__(self, credentials=None):
    self._state_handler_cache = {}
    self._lock = threading.Lock()
    self._throwing_state_handler = ThrowingStateHandler()
    self._credentials = credentials

  def create_state_handler(self, api_service_descriptor):
    if not api_service_descriptor:
      return self._throwing_state_handler
    url = api_service_descriptor.url
    if url not in self._state_handler_cache:
      with self._lock:
        if url not in self._state_handler_cache:
          # 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.
          options = [('grpc.max_receive_message_length', -1),
                     ('grpc.max_send_message_length', -1)]
          if self._credentials is None:
            logging.info('Creating insecure state channel for %s.', url)
            grpc_channel = GRPCChannelFactory.insecure_channel(
                url, options=options)
          else:
            logging.info('Creating secure state channel for %s.', url)
            grpc_channel = GRPCChannelFactory.secure_channel(
                url, self._credentials, options=options)
          logging.info('State channel established.')
          # Add workerId to the grpc channel
          grpc_channel = grpc.intercept_channel(grpc_channel,
                                                WorkerIdInterceptor())
          self._state_handler_cache[url] = GrpcStateHandler(
              beam_fn_api_pb2_grpc.BeamFnStateStub(grpc_channel))
    return self._state_handler_cache[url]

  def close(self):
    logging.info('Closing all cached gRPC state handlers.')
    for _, state_handler in self._state_handler_cache.items():
      state_handler.done()
    self._state_handler_cache.clear()


class ThrowingStateHandler(object):
  """A state handler that errors on any requests."""

  def blocking_get(self, state_key, instruction_reference):
    raise RuntimeError(
        'Unable to handle state requests for ProcessBundleDescriptor without '
        'out state ApiServiceDescriptor for instruction %s and state key %s.'
        % (state_key, instruction_reference))

  def blocking_append(self, state_key, data, instruction_reference):
    raise RuntimeError(
        'Unable to handle state requests for ProcessBundleDescriptor without '
        'out state ApiServiceDescriptor for instruction %s and state key %s.'
        % (state_key, instruction_reference))

  def blocking_clear(self, state_key, instruction_reference):
    raise RuntimeError(
        'Unable to handle state requests for ProcessBundleDescriptor without '
        'out state ApiServiceDescriptor for instruction %s and state key %s.'
        % (state_key, instruction_reference))


class GrpcStateHandler(object):

  _DONE = object()

  def __init__(self, state_stub):
    self._lock = threading.Lock()
    self._state_stub = state_stub
    self._requests = queue.Queue()
    self._responses_by_id = {}
    self._last_id = 0
    self._exc_info = None
    self._context = threading.local()
    self.start()

  @contextlib.contextmanager
  def process_instruction_id(self, bundle_id):
    if getattr(self._context, 'process_instruction_id', None) is not None:
      raise RuntimeError(
          'Already bound to %r' % self._context.process_instruction_id)
    self._context.process_instruction_id = bundle_id
    try:
      yield
    finally:
      self._context.process_instruction_id = None

  def start(self):
    self._done = False

    def request_iter():
      while True:
        request = self._requests.get()
        if request is self._DONE or self._done:
          break
        yield request

    responses = self._state_stub.State(request_iter())

    def pull_responses():
      try:
        for response in responses:
          self._responses_by_id[response.id].set(response)
          if self._done:
            break
      except:  # pylint: disable=bare-except
        self._exc_info = sys.exc_info()
        raise

    reader = threading.Thread(target=pull_responses, name='read_state')
    reader.daemon = True
    reader.start()

  def done(self):
    self._done = True
    self._requests.put(self._DONE)

  def blocking_get(self, state_key, continuation_token=None):
    response = self._blocking_request(
        beam_fn_api_pb2.StateRequest(
            state_key=state_key,
            get=beam_fn_api_pb2.StateGetRequest(
                continuation_token=continuation_token)))
    return response.get.data, response.get.continuation_token

  def blocking_append(self, state_key, data):
    self._blocking_request(
        beam_fn_api_pb2.StateRequest(
            state_key=state_key,
            append=beam_fn_api_pb2.StateAppendRequest(data=data)))

  def blocking_clear(self, state_key):
    self._blocking_request(
        beam_fn_api_pb2.StateRequest(
            state_key=state_key,
            clear=beam_fn_api_pb2.StateClearRequest()))

  def _blocking_request(self, request):
    request.id = self._next_id()
    request.instruction_reference = self._context.process_instruction_id
    self._responses_by_id[request.id] = future = _Future()
    self._requests.put(request)
    while not future.wait(timeout=1):
      if self._exc_info:
        t, v, tb = self._exc_info
        raise_(t, v, tb)
      elif self._done:
        raise RuntimeError()
    del self._responses_by_id[request.id]
    response = future.get()
    if response.error:
      raise RuntimeError(response.error)
    else:
      return response

  def _next_id(self):
    self._last_id += 1
    return str(self._last_id)


class _Future(object):
  """A simple future object to implement blocking requests.
  """

  def __init__(self):
    self._event = threading.Event()

  def wait(self, timeout=None):
    return self._event.wait(timeout)

  def get(self, timeout=None):
    if self.wait(timeout):
      return self._value
    else:
      raise LookupError()

  def set(self, value):
    self._value = value
    self._event.set()