#
# 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.
#
import functools
import logging
import os
import Queue as queue
import subprocess
import threading
import time
import traceback
import uuid
from concurrent import futures

import grpc
from google.protobuf import text_format

from apache_beam.portability.api import beam_fn_api_pb2_grpc
from apache_beam.portability.api import beam_job_api_pb2
from apache_beam.portability.api import beam_job_api_pb2_grpc
from apache_beam.portability.api import endpoints_pb2
from apache_beam.runners.portability import fn_api_runner

TERMINAL_STATES = [
    beam_job_api_pb2.JobState.DONE,
    beam_job_api_pb2.JobState.STOPPED,
    beam_job_api_pb2.JobState.FAILED,
    beam_job_api_pb2.JobState.CANCELLED,
]


class LocalJobServicer(beam_job_api_pb2_grpc.JobServiceServicer):
  """
    Experimental: No backward compatibility guaranteed.
    Servicer for the Beam Job API.

    Manages one or more pipelines, possibly concurrently.

    This JobService uses a basic local implementation of runner to run the job.
    This JobService is not capable of managing job on remote clusters.

    By default, this JobService executes the job in process but still uses GRPC
    to communicate pipeline and worker state.  It can also be configured to use
    inline calls rather than GRPC (for speed) or launch completely separate
    subprocesses for the runner and worker(s).
    """

  def __init__(self, worker_command_line=None, use_grpc=True):
    self._worker_command_line = worker_command_line
    self._use_grpc = use_grpc or bool(worker_command_line)
    self._jobs = {}

  def start_grpc_server(self, port=0):
    self._server = grpc.server(futures.ThreadPoolExecutor(max_workers=3))
    port = self._server.add_insecure_port('localhost:%d' % port)
    beam_job_api_pb2_grpc.add_JobServiceServicer_to_server(self, self._server)
    self._server.start()
    logging.info('Grpc server started on port %s', port)
    return port

  def Prepare(self, request, context=None):
    # For now, just use the job name as the job id.
    logging.debug('Got Prepare request.')
    preparation_id = '%s-%s' % (request.job_name, uuid.uuid4())
    if self._worker_command_line:
      sdk_harness_factory = functools.partial(SubprocessSdkWorker,
                                              self._worker_command_line)
    else:
      sdk_harness_factory = None
    self._jobs[preparation_id] = BeamJob(
        preparation_id,
        request.pipeline_options,
        request.pipeline,
        use_grpc=self._use_grpc,
        sdk_harness_factory=sdk_harness_factory)
    logging.debug("Prepared job '%s' as '%s'", request.job_name, preparation_id)
    # TODO(angoenka): Pass an appropriate staging_session_token. The token can
    # be obtained in PutArtifactResponse from JobService
    return beam_job_api_pb2.PrepareJobResponse(
        preparation_id=preparation_id, staging_session_token='token')

  def Run(self, request, context=None):
    job_id = request.preparation_id
    logging.debug("Runing job '%s'", job_id)
    self._jobs[job_id].start()
    return beam_job_api_pb2.RunJobResponse(job_id=job_id)

  def GetState(self, request, context=None):
    return beam_job_api_pb2.GetJobStateResponse(
        state=self._jobs[request.job_id].state)

  def Cancel(self, request, context=None):
    self._jobs[request.job_id].cancel()
    return beam_job_api_pb2.CancelJobRequest(
        state=self._jobs[request.job_id].state)

  def GetStateStream(self, request, context=None):
    job = self._jobs[request.job_id]
    state_queue = queue.Queue()
    job.add_state_change_callback(lambda state: state_queue.put(state))
    try:
      current_state = state_queue.get()
    except queue.Empty:
      current_state = job.state
    yield beam_job_api_pb2.GetJobStateResponse(state=current_state)
    while current_state not in TERMINAL_STATES:
      current_state = state_queue.get(block=True)
      yield beam_job_api_pb2.GetJobStateResponse(state=current_state)

  def GetMessageStream(self, request, context=None):
    job = self._jobs[request.job_id]
    current_state = job.state
    while current_state not in TERMINAL_STATES:
      msg = job.log_queue.get(block=True)
      yield msg
      if msg.HasField('state_response'):
        current_state = msg.state_response.state
    try:
      while True:
        yield job.log_queue.get(block=False)
    except queue.Empty:
      pass


class SubprocessSdkWorker(object):
  """Manages a SDK worker implemented as a subprocess communicating over grpc.
    """

  def __init__(self, worker_command_line, control_address):
    self._worker_command_line = worker_command_line
    self._control_address = control_address

  def run(self):
    logging_server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    logging_port = logging_server.add_insecure_port('[::]:0')
    logging_server.start()
    logging_servicer = BeamFnLoggingServicer()
    beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server(
        logging_servicer, logging_server)
    logging_descriptor = text_format.MessageToString(
        endpoints_pb2.ApiServiceDescriptor(url='localhost:%s' % logging_port))

    control_descriptor = text_format.MessageToString(
        endpoints_pb2.ApiServiceDescriptor(url=self._control_address))

    p = subprocess.Popen(
        self._worker_command_line,
        shell=True,
        env=dict(
            os.environ,
            CONTROL_API_SERVICE_DESCRIPTOR=control_descriptor,
            LOGGING_API_SERVICE_DESCRIPTOR=logging_descriptor))
    try:
      p.wait()
      if p.returncode:
        raise RuntimeError(
            'Worker subprocess exited with return code %s' % p.returncode)
    finally:
      if p.poll() is None:
        p.kill()
      logging_server.stop(0)


class BeamJob(threading.Thread):
  """This class handles running and managing a single pipeline.

    The current state of the pipeline is available as self.state.
    """

  def __init__(self,
               job_id,
               pipeline_options,
               pipeline_proto,
               use_grpc=True,
               sdk_harness_factory=None):
    super(BeamJob, self).__init__()
    self._job_id = job_id
    self._pipeline_options = pipeline_options
    self._pipeline_proto = pipeline_proto
    self._use_grpc = use_grpc
    self._sdk_harness_factory = sdk_harness_factory
    self._log_queue = queue.Queue()
    self._state_change_callbacks = [
        lambda new_state: self._log_queue.put(
            beam_job_api_pb2.JobMessagesResponse(
                state_response=
                beam_job_api_pb2.GetJobStateResponse(state=new_state)))
    ]
    self._state = None
    self.state = beam_job_api_pb2.JobState.STARTING
    self.daemon = True

  def add_state_change_callback(self, f):
    self._state_change_callbacks.append(f)

  @property
  def log_queue(self):
    return self._log_queue

  @property
  def state(self):
    return self._state

  @state.setter
  def state(self, new_state):
    for state_change_callback in self._state_change_callbacks:
      state_change_callback(new_state)
    self._state = new_state

  def run(self):
    with JobLogHandler(self._log_queue):
      try:
        fn_api_runner.FnApiRunner(
            use_grpc=self._use_grpc,
            sdk_harness_factory=self._sdk_harness_factory).run_via_runner_api(
                self._pipeline_proto)
        logging.info('Successfully completed job.')
        self.state = beam_job_api_pb2.JobState.DONE
      except:  # pylint: disable=bare-except
        logging.exception('Error running pipeline.')
        traceback.print_exc()
        self.state = beam_job_api_pb2.JobState.FAILED
        raise

  def cancel(self):
    if self.state not in TERMINAL_STATES:
      self.state = beam_job_api_pb2.JobState.CANCELLING
      # TODO(robertwb): Actually cancel...
      self.state = beam_job_api_pb2.JobState.CANCELLED


class BeamFnLoggingServicer(beam_fn_api_pb2_grpc.BeamFnLoggingServicer):

  def Logging(self, log_bundles, context=None):
    for log_bundle in log_bundles:
      for log_entry in log_bundle.log_entries:
        logging.info('Worker: %s', str(log_entry).replace('\n', ' '))
    return iter([])


class JobLogHandler(logging.Handler):
  """Captures logs to be returned via the Beam Job API.

    Enabled via the with statement."""

  # Mapping from logging levels to LogEntry levels.
  LOG_LEVEL_MAP = {
      logging.FATAL: beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR,
      logging.ERROR: beam_job_api_pb2.JobMessage.JOB_MESSAGE_ERROR,
      logging.WARNING: beam_job_api_pb2.JobMessage.JOB_MESSAGE_WARNING,
      logging.INFO: beam_job_api_pb2.JobMessage.JOB_MESSAGE_BASIC,
      logging.DEBUG: beam_job_api_pb2.JobMessage.JOB_MESSAGE_DEBUG,
  }

  def __init__(self, message_queue):
    super(JobLogHandler, self).__init__()
    self._message_queue = message_queue
    self._last_id = 0
    self._logged_thread = None

  def __enter__(self):
    # Remember the current thread to demultiplex the logs of concurrently
    # running pipelines (as Python log handlers are global).
    self._logged_thread = threading.current_thread()
    logging.getLogger().addHandler(self)

  def __exit__(self, *args):
    self._logged_thread = None
    self.close()

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

  def emit(self, record):
    if self._logged_thread is threading.current_thread():
      self._message_queue.put(
          beam_job_api_pb2.JobMessagesResponse(
              message_response=beam_job_api_pb2.JobMessage(
                  message_id=self._next_id(),
                  time=time.strftime('%Y-%m-%d %H:%M:%S.',
                                     time.localtime(record.created)),
                  importance=self.LOG_LEVEL_MAP[record.levelno],
                  message_text=self.format(record))))
