"""ZeroMQ implementations of the Plaso queue interface."""

import abc
import errno
import queue
import threading
import time

import zmq

from plaso.engine import logger
from plaso.lib import errors
from plaso.multi_process import plaso_queue


# pylint: disable=no-member

class ZeroMQQueue(plaso_queue.Queue):
  """Interface for a ZeroMQ backed queue.

  Attributes:
    name (str): name to identify the queue.
    port (int): TCP port that the queue is connected or bound to. If the queue
        is not yet bound or connected to a port, this value will be None.
    timeout_seconds (int): number of seconds that calls to PopItem and PushItem
        may block for, before returning queue.QueueEmpty.
  """

  _SOCKET_ADDRESS = 'tcp://127.0.0.1'
  _SOCKET_TYPE = None

  _ZMQ_SOCKET_SEND_TIMEOUT_MILLISECONDS = 1500
  _ZMQ_SOCKET_RECEIVE_TIMEOUT_MILLISECONDS = 1500

  SOCKET_CONNECTION_BIND = 1
  SOCKET_CONNECTION_CONNECT = 2
  SOCKET_CONNECTION_TYPE = None

  def __init__(
      self, delay_open=True, linger_seconds=10, maximum_items=1000,
      name='Unnamed', port=None, timeout_seconds=5):
    """Initializes a ZeroMQ backed queue.

    Args:
      delay_open (Optional[bool]): whether a ZeroMQ socket should be created
          the first time the queue is pushed to or popped from, rather than at
          queue object initialization. This is useful if a queue needs to be
          passed to a child process from a parent process.
      linger_seconds (Optional[int]): number of seconds that the underlying
          ZeroMQ socket can remain open after the queue has been closed,
          to allow queued items to be transferred to other ZeroMQ sockets.
      maximum_items (Optional[int]): maximum number of items to queue on the
          ZeroMQ socket. ZeroMQ refers to this value as "high water mark" or
          "hwm". Note that this limit only applies at one "end" of the queue.
          The default of 1000 is the ZeroMQ default value.
      name (Optional[str]): Optional name to identify the queue.
      port (Optional[int]): The TCP port to use for the queue. The default is
          None, which indicates that the queue should choose a random port to
          bind to.
      timeout_seconds (Optional[int]): number of seconds that calls to PopItem
          and PushItem may block for, before returning queue.QueueEmpty.

    Raises:
      ValueError: if the queue is configured to connect to an endpoint,
          but no port is specified.
    """
    if (self.SOCKET_CONNECTION_TYPE == self.SOCKET_CONNECTION_CONNECT
        and not port):
      raise ValueError('No port specified to connect to.')

    super().__init__()
    self._closed_event = None
    self._high_water_mark = maximum_items
    self._linger_seconds = linger_seconds
    self._terminate_event = None
    self._zmq_context = None
    self._zmq_socket = None

    self.name = name
    self.port = port
    self.timeout_seconds = timeout_seconds

    if not delay_open:
      self._CreateZMQSocket()

  def _SendItem(self, zmq_socket, item, block=True):
    """Attempts to send an item to a ZeroMQ socket.

    Args:
      zmq_socket (zmq.Socket): used to the send the item.
      item (object): sent on the queue. Will be pickled prior to sending.
      block (Optional[bool]): whether the push should be performed in blocking
          or non-blocking mode.

    Returns:
      bool: whether the item was sent successfully.
    """
    try:
      logger.debug(f'[{self.name:s}] Sending item.')
      if block:
        zmq_socket.send_pyobj(item)
      else:
        zmq_socket.send_pyobj(item, zmq.DONTWAIT)
      logger.debug(f'[{self.name:s}] Sent item.')
      return True

    except zmq.error.Again:
      logger.debug(f'[{self.name:s}] Could not send an item.')

    except zmq.error.ZMQError as exception:
      if exception.errno == errno.EINTR:
        logger.error(f'[{self.name:s}] ZMQ syscall interrupted.')

    return False

  def _ReceiveItemOnActivity(self, zmq_socket):
    """Attempts to receive an item from a ZeroMQ socket.

    Args:
      zmq_socket (zmq.Socket): used to the receive the item.

    Returns:
      object: item from the socket.

    Raises:
      QueueEmpty: if no item could be received within the timeout.
      zmq.error.ZMQError: if an error occurs in ZeroMQ
    """
    events = zmq_socket.poll(
        self._ZMQ_SOCKET_RECEIVE_TIMEOUT_MILLISECONDS)
    if events:
      try:
        received_object = self._zmq_socket.recv_pyobj()
        return received_object

      except zmq.error.Again:
        logger.error(f'[{self.name:s}] Failed to receive item in time.')
        raise

      except zmq.error.ZMQError as exception:
        if exception.errno == errno.EINTR:
          logger.error(f'[{self.name:s}] ZMQ syscall interrupted, aborting.')
        raise

    raise errors.QueueEmpty

  def _SetSocketTimeouts(self):
    """Sets the timeouts for socket send and receive."""
    # Note that timeout must be an integer value. If timeout is a float
    # it appears that zmq will not enforce the timeout.
    timeout = int(self.timeout_seconds * 1000)
    receive_timeout = min(
        self._ZMQ_SOCKET_RECEIVE_TIMEOUT_MILLISECONDS, timeout)
    send_timeout = min(self._ZMQ_SOCKET_SEND_TIMEOUT_MILLISECONDS, timeout)

    self._zmq_socket.setsockopt(zmq.RCVTIMEO, receive_timeout)
    self._zmq_socket.setsockopt(zmq.SNDTIMEO, send_timeout)

  def _SetSocketHighWaterMark(self):
    """Sets the high water mark for the socket.

    This number is the maximum number of items that will be queued in the socket
    on this end of the queue.
    """
    self._zmq_socket.hwm = self._high_water_mark

  def _CreateZMQSocket(self):
    """Creates a ZeroMQ socket."""
    logger.debug(f'[{self.name:s}] Creating socket.')

    if not self._zmq_context:
      self._zmq_context = zmq.Context()  # pylint: disable=abstract-class-instantiated

    # The terminate and close threading events need to be created when the
    # socket is opened. Threading events are unpickleable objects and cannot
    # passed in multiprocessing on Windows.

    if not self._terminate_event:
      self._terminate_event = threading.Event()

    if not self._closed_event:
      self._closed_event = threading.Event()

    if self._zmq_socket:
      logger.debug(f'[{self.name:s}] Closing old socket.')
      self._zmq_socket.close()
      self._zmq_socket = None

    self._zmq_socket = self._zmq_context.socket(self._SOCKET_TYPE)
    self._SetSocketTimeouts()
    self._SetSocketHighWaterMark()

    if self.port:
      address = f'{self._SOCKET_ADDRESS:s}:{self.port:d}'
      if self.SOCKET_CONNECTION_TYPE == self.SOCKET_CONNECTION_CONNECT:
        self._zmq_socket.connect(address)
        logger.debug(f'[{self.name:s}] Connected to {address:s}')
      else:
        self._zmq_socket.bind(address)
        logger.debug(f'[{self.name:s}] Bound to port {address:s}')
    else:
      self.port = self._zmq_socket.bind_to_random_port(self._SOCKET_ADDRESS)
      logger.debug(f'[{self.name:s}] Bound to port {self.port:d}')

  def Open(self):
    """Opens this queue, causing the creation of a ZeroMQ socket.

    Raises:
      QueueAlreadyStarted: if the queue is already started, and a socket already
          exists.
    """
    if self._zmq_socket:
      raise errors.QueueAlreadyStarted()

    self._CreateZMQSocket()

  def Close(self, abort=False):
    """Closes the queue.

    Args:
      abort (Optional[bool]): whether the Close is the result of an abort
          condition. If True, queue contents may be lost.

    Raises:
      QueueAlreadyClosed: if the queue is not started, or has already been
          closed.
      RuntimeError: if closed or terminate event is missing.
    """
    if not self._closed_event or not self._terminate_event:
      raise RuntimeError('Missing closed or terminate event.')

    if not abort and self._closed_event.is_set():
      raise errors.QueueAlreadyClosed()

    self._closed_event.set()

    if abort:
      if not self._closed_event.is_set():
        logger.warning(f'[{self.name:s}] Queue aborting, contents may be lost.')

      self._linger_seconds = 0

      # We can't determine whether a there might be an operation being performed
      # on the socket in a separate method or thread, so we'll signal that any
      # such operation should cease.
      self._terminate_event.set()

    else:
      logger.debug((
          f'[{self.name:s}] Queue closing, will linger for up to '
          f'{self._linger_seconds:d} seconds.'))

  def IsBound(self):
    """Checks if the queue is bound to a port."""
    return (self.SOCKET_CONNECTION_TYPE == self.SOCKET_CONNECTION_BIND and
            self.port is not None)

  def IsConnected(self):
    """Checks if the queue is connected to a port."""
    return (self.SOCKET_CONNECTION_TYPE == self.SOCKET_CONNECTION_CONNECT and
            self.port is not None)

  def IsEmpty(self):
    """Checks if the queue is empty.

    ZeroMQ queues don't have a concept of "empty" - there could always be
    messages on the queue that a producer or consumer is unaware of. Thus,
    the queue is never empty, so we return False. Note that it is possible that
    a queue is unable to pop an item from a queue within a timeout, which will
    cause PopItem to raise a QueueEmpty exception, but this is a different
    condition.

    Returns:
      bool: False, to indicate the the queue isn't empty.
    """
    return False

  @abc.abstractmethod
  def PushItem(self, item, block=True):
    """Pushes an item on to the queue.

    Args:
      item (object): item to push on the queue.
      block (Optional[bool]): whether the push should be performed in blocking
          or non-blocking mode.

    Raises:
      QueueAlreadyClosed: if the queue is closed.
    """

  # pylint: disable=redundant-returns-doc
  @abc.abstractmethod
  def PopItem(self):
    """Pops an item off the queue.

    Returns:
      object: item from the queue.

    Raises:
      QueueEmpty: if the queue is empty, and no item could be popped within the
          queue timeout.
    """


class ZeroMQPullQueue(ZeroMQQueue):
  """Parent class for Plaso queues backed by ZeroMQ PULL sockets.

  This class should not be instantiated directly, a subclass should be
  instantiated instead.

  Instances of this class or subclasses may only be used to pop items, not to
  push.
  """

  _SOCKET_TYPE = zmq.PULL

  def PopItem(self):
    """Pops an item off the queue.

    If no ZeroMQ socket has been created, one will be created the first
    time this method is called.

    Returns:
      object: item from the queue.

    Raises:
      KeyboardInterrupt: if the process is sent a KeyboardInterrupt while
          popping an item.
      QueueEmpty: if the queue is empty, and no item could be popped within the
          queue timeout.
      RuntimeError: if closed or terminate event is missing.
      zmq.error.ZMQError: if a ZeroMQ error occurs.
    """
    if not self._zmq_socket:
      self._CreateZMQSocket()

    if not self._closed_event or not self._terminate_event:
      raise RuntimeError('Missing closed or terminate event.')

    logger.debug(
        f'Pop on {self.name:s} queue, port {self.port:d}')

    last_retry_timestamp = time.time() + self.timeout_seconds
    while not self._closed_event.is_set() or not self._terminate_event.is_set():
      try:
        return self._ReceiveItemOnActivity(self._zmq_socket)

      except errors.QueueEmpty:
        if time.time() > last_retry_timestamp:
          raise

      except KeyboardInterrupt:
        self.Close(abort=True)
        raise

    return None

  def PushItem(self, item, block=True):
    """Pushes an item on to the queue.

    Provided for compatibility with the API, but doesn't actually work.

    Args:
      item (object): item to push on the queue.
      block (Optional[bool]): whether the push should be performed in blocking
          or non-blocking mode.

    Raises:
      WrongQueueType: As Push is not supported this queue.
    """
    raise errors.WrongQueueType()


class ZeroMQPullConnectQueue(ZeroMQPullQueue):
  """A Plaso queue backed by a ZeroMQ PULL socket that connects to a port.

  This queue may only be used to pop items, not to push.
  """
  SOCKET_CONNECTION_TYPE = ZeroMQQueue.SOCKET_CONNECTION_CONNECT


class ZeroMQPushQueue(ZeroMQQueue):
  """Parent class for Plaso queues backed by ZeroMQ PUSH sockets.

  This class should not be instantiated directly, a subclass should be
  instantiated instead.

  Instances of this class or subclasses may only be used to push items, not to
  pop.
  """

  _SOCKET_TYPE = zmq.PUSH

  def PopItem(self):
    """Pops an item of the queue.

    Provided for compatibility with the API, but doesn't actually work.

    Raises:
      WrongQueueType: As Pull is not supported this queue.
    """
    raise errors.WrongQueueType()

  def PushItem(self, item, block=True):
    """Push an item on to the queue.

    If no ZeroMQ socket has been created, one will be created the first time
    this method is called.

    Args:
      item (object): item to push on the queue.
      block (Optional[bool]): whether the push should be performed in blocking
          or non-blocking mode.

    Raises:
      KeyboardInterrupt: if the process is sent a KeyboardInterrupt while
          pushing an item.
      QueueFull: if it was not possible to push the item to the queue
          within the timeout.
      RuntimeError: if terminate event is missing.
      zmq.error.ZMQError: if a ZeroMQ specific error occurs.
    """
    if not self._zmq_socket:
      self._CreateZMQSocket()

    if not self._terminate_event:
      raise RuntimeError('Missing terminate event.')

    logger.debug(f'Push on {self.name:s} queue, port {self.port:d}')

    last_retry_timestamp = time.time() + self.timeout_seconds
    while not self._terminate_event.is_set():
      try:
        send_successful = self._SendItem(self._zmq_socket, item, block)
        if send_successful:
          break

        if time.time() > last_retry_timestamp:
          logger.error(f'{self.name:s} unable to push item, raising.')
          raise errors.QueueFull

      except KeyboardInterrupt:
        self.Close(abort=True)
        raise


class ZeroMQPushBindQueue(ZeroMQPushQueue):
  """A Plaso queue backed by a ZeroMQ PUSH socket that binds to a port.

  This queue may only be used to push items, not to pop.
  """
  SOCKET_CONNECTION_TYPE = ZeroMQQueue.SOCKET_CONNECTION_BIND


class ZeroMQRequestQueue(ZeroMQQueue):
  """Parent class for Plaso queues backed by ZeroMQ REQ sockets.

  This class should not be instantiated directly, a subclass should be
  instantiated instead.

  Instances of this class or subclasses may only be used to pop items, not to
  push.
  """

  _SOCKET_TYPE = zmq.REQ

  def PopItem(self):
    """Pops an item off the queue.

    If no ZeroMQ socket has been created, one will be created the first
    time this method is called.

    Returns:
      object: item from the queue.

    Raises:
      KeyboardInterrupt: if the process is sent a KeyboardInterrupt while
          popping an item.
      QueueEmpty: if the queue is empty, and no item could be popped within the
          queue timeout.
      RuntimeError: if terminate event is missing.
      zmq.error.ZMQError: if an error occurs in ZeroMQ.
    """
    if not self._zmq_socket:
      self._CreateZMQSocket()

    if not self._terminate_event:
      raise RuntimeError('Missing terminate event.')

    logger.debug(f'Pop on {self.name:s} queue, port {self.port:d}')

    last_retry_time = time.time() + self.timeout_seconds
    while not self._terminate_event.is_set():
      try:
        self._zmq_socket.send_pyobj(None)
        break

      except zmq.error.Again:
        # The existing socket is now out of sync, so we need to open a new one.
        self._CreateZMQSocket()
        if time.time() > last_retry_time:
          logger.warning(f'{self.name:s} timeout requesting item.')
          raise errors.QueueEmpty

        continue

    while not self._terminate_event.is_set():
      try:
        return self._ReceiveItemOnActivity(self._zmq_socket)
      except errors.QueueEmpty:
        continue

      except KeyboardInterrupt:
        self.Close(abort=True)
        raise

    return None

  def PushItem(self, item, block=True):
    """Pushes an item on to the queue.

    Provided for compatibility with the API, but doesn't actually work.

    Args:
      item (object): item to push on the queue.
      block (Optional[bool]): whether the push should be performed in blocking
          or non-blocking mode.

    Raises:
      WrongQueueType: As Push is not supported this queue.
    """
    raise errors.WrongQueueType


class ZeroMQRequestConnectQueue(ZeroMQRequestQueue):
  """A Plaso queue backed by a ZeroMQ REQ socket that connects to a port.

  This queue may only be used to pop items, not to push.
  """
  SOCKET_CONNECTION_TYPE = ZeroMQQueue.SOCKET_CONNECTION_CONNECT


class ZeroMQBufferedQueue(ZeroMQQueue):
  """Parent class for buffered Plaso queues.

  Buffered queues use a regular Python queue to store items that are pushed or
  popped from the queue without blocking on underlying ZeroMQ operations.

  This class should not be instantiated directly, a subclass should be
  instantiated instead.
  """

  def __init__(
      self, buffer_timeout_seconds=2, buffer_max_size=10000, delay_open=True,
      linger_seconds=10, maximum_items=1000, name='Unnamed', port=None,
      timeout_seconds=5):
    """Initializes a buffered, ZeroMQ backed queue.

    Args:
      buffer_max_size (Optional[int]): maximum number of items to store in
          the buffer, before or after they are sent/received via ZeroMQ.
      buffer_timeout_seconds(Optional[int]): number of seconds to wait when
          doing a put or get to/from the internal buffer.
      delay_open (Optional[bool]): whether a ZeroMQ socket should be created
          the first time the queue is pushed to or popped from, rather than at
          queue object initialization. This is useful if a queue needs to be
          passed to a child process from a parent process.
      linger_seconds (Optional[int]): number of seconds that the underlying
          ZeroMQ socket can remain open after the queue object has been closed,
          to allow queued items to be transferred to other ZeroMQ sockets.
      maximum_items (Optional[int]): maximum number of items to queue on the
          ZeroMQ socket. ZeroMQ refers to this value as "high water mark" or
          "hwm". Note that this limit only applies at one "end" of the queue.
          The default of 1000 is the ZeroMQ default value.
      name (Optional[str]): name to identify the queue.
      port (Optional[int]): The TCP port to use for the queue. None indicates
          that the queue should choose a random port to bind to.
      timeout_seconds (Optional[int]): number of seconds that calls to PopItem
          and PushItem may block for, before returning queue.QueueEmpty.
    """
    self._buffer_timeout_seconds = buffer_timeout_seconds
    self._queue = queue.Queue(maxsize=buffer_max_size)
    self._zmq_thread = None

    # We need to set up the internal buffer queue before we call super, so that
    # if the call to super opens the ZMQSocket, the backing thread will work.
    super().__init__(
        delay_open=delay_open, linger_seconds=linger_seconds,
        maximum_items=maximum_items, name=name, port=port,
        timeout_seconds=timeout_seconds)

  def _CreateZMQSocket(self):
    """Creates a ZeroMQ socket as well as a regular queue and a thread."""
    super()._CreateZMQSocket()
    if not self._zmq_thread:
      thread_name = f'{self.name:s}_zmq_responder'
      self._zmq_thread = threading.Thread(
          target=self._ZeroMQResponder, args=[self._queue], name=thread_name)
      self._zmq_thread.start()

  @abc.abstractmethod
  def _ZeroMQResponder(self, source_queue):
    """Listens for requests and replies to clients.

    Args:
      source_queue (queue.queue): queue to to pull items from.
    """

  def Close(self, abort=False):
    """Closes the queue.

    Args:
      abort (Optional[bool]): whether the Close is the result of an abort
          condition. If True, queue contents may be lost.

    Raises:
      QueueAlreadyClosed: if the queue is not started, or has already been
          closed.
      RuntimeError: if closed or terminate event is missing.
    """
    if not self._closed_event or not self._terminate_event:
      raise RuntimeError('Missing closed or terminate event.')

    if not abort and self._closed_event.is_set():
      raise errors.QueueAlreadyClosed()

    self._closed_event.set()

    if abort:
      if not self._closed_event.is_set():
        logger.warning(f'{self.name:s} queue aborting. Contents may be lost.')

      # We can't determine whether a there might be an operation being performed
      # on the socket in a separate method or thread, so we'll signal that any
      # such operation should cease.
      self._terminate_event.set()

      self._linger_seconds = 0

      if self._zmq_thread:
        logger.debug(f'[{self.name:s}] Waiting for thread to exit.')
        self._zmq_thread.join(timeout=self.timeout_seconds)
        if self._zmq_thread.is_alive():
          logger.error((
              f'[{self.name:s}] ZMQ responder thread did not exit within '
              f'timeout.'))
    else:
      logger.debug((
          f'[{self.name:s}] Queue closing, will linger for up to '
          f'{self._linger_seconds:d} seconds.'))

  def Empty(self):
    """Removes all items from the internal buffer."""
    try:
      while True:
        self._queue.get(False)
    except queue.Empty:
      pass


class ZeroMQBufferedReplyQueue(ZeroMQBufferedQueue):
  """Parent class for buffered Plaso queues backed by ZeroMQ REP sockets.

  This class should not be instantiated directly, a subclass should be
  instantiated instead.

  Instances of this class or subclasses may only be used to push items, not to
  pop.
  """

  _ZMQ_SOCKET_RECEIVE_TIMEOUT_MILLISECONDS = 4000
  _ZMQ_SOCKET_SEND_TIMEOUT_MILLISECONDS = 2000

  _SOCKET_TYPE = zmq.REP

  def _ZeroMQResponder(self, source_queue):
    """Listens for requests and replies to clients.

    Args:
      source_queue (queue.Queue): queue to use to pull items from.

    Raises:
      RuntimeError: if closed or terminate event is missing.
    """
    if not self._closed_event or not self._terminate_event:
      raise RuntimeError('Missing closed or terminate event.')

    logger.debug(f'[{self.name:s}] responder thread started.')

    item = None
    while not self._terminate_event.is_set():
      if not item:
        try:
          if self._closed_event.is_set():
            item = source_queue.get_nowait()
          else:
            item = source_queue.get(True, self._buffer_timeout_seconds)

        except queue.Empty:
          if self._closed_event.is_set():
            break

          continue

      try:
        # We need to receive a request before we can reply with the item.
        self._ReceiveItemOnActivity(self._zmq_socket)

      except errors.QueueEmpty:
        if self._closed_event.is_set() and self._queue.empty():
          break

        continue

      sent_successfully = self._SendItem(self._zmq_socket, item)
      item = None
      if not sent_successfully:
        logger.error(f'{self.name:s}] Unable to send item.')
        break

    logger.info(f'[{self.name:s}] Responder exiting.')
    self._zmq_socket.close(self._linger_seconds)

  def PopItem(self):
    """Pops an item of the queue.

    Provided for compatibility with the API, but doesn't actually work.

    Raises:
      WrongQueueType: As Pop is not supported by this queue.
    """
    raise errors.WrongQueueType()

  def PushItem(self, item, block=True):
    """Push an item on to the queue.

    If no ZeroMQ socket has been created, one will be created the first time
    this method is called.

    Args:
      item (object): item to push on the queue.
      block (Optional[bool]): whether the push should be performed in blocking
          or non-blocking mode.

    Raises:
      QueueAlreadyClosed: if the queue is closed.
      QueueFull: if the internal buffer was full and it was not possible to
          push the item to the buffer within the timeout.
      RuntimeError: if closed event is missing.
    """
    if not self._closed_event:
      raise RuntimeError('Missing closed event.')

    if self._closed_event.is_set():
      raise errors.QueueAlreadyClosed()

    if not self._zmq_socket:
      self._CreateZMQSocket()

    try:
      if block:
        self._queue.put(item, timeout=self.timeout_seconds)
      else:
        self._queue.put(item, block=False)
    except queue.Full as exception:
      raise errors.QueueFull(exception)


class ZeroMQBufferedReplyBindQueue(ZeroMQBufferedReplyQueue):
  """A Plaso queue backed by a ZeroMQ REP socket that binds to a port.

  This queue may only be used to pop items, not to push.
  """
  SOCKET_CONNECTION_TYPE = ZeroMQQueue.SOCKET_CONNECTION_BIND
