"""The multi-process processing engine."""

import abc
import ctypes
import os
import signal
import sys
import threading
import time

from plaso.engine import engine
from plaso.engine import process_info
from plaso.lib import definitions
from plaso.multi_process import logger
from plaso.multi_process import plaso_xmlrpc


class MultiProcessEngine(engine.BaseEngine):
  """Multi-process engine base.

  This class contains functionality to:
  * monitor and manage worker processes;
  * retrieve a process status information via RPC;
  * manage the status update thread.
  """

  # Note that on average Windows seems to require a longer wait.
  _RPC_SERVER_TIMEOUT = 8.0
  _MAXIMUM_RPC_ERRORS = 10
  # Maximum number of attempts to try to start a replacement worker process.
  _MAXIMUM_REPLACEMENT_RETRIES = 3
  # Number of seconds to wait between attempts to start a replacement worker
  # process
  _REPLACEMENT_WORKER_RETRY_DELAY = 1

  _PROCESS_JOIN_TIMEOUT = 5.0

  _ZEROMQ_NO_WORKER_REQUEST_TIME_SECONDS = 300

  def __init__(self):
    """Initializes a multi-process engine."""
    super().__init__()
    self._debug_output = False
    self._name = 'Main'
    self._last_worker_number = 0
    self._log_filename = None
    self._pid = os.getpid()
    self._process_information = process_info.ProcessInfo(self._pid)
    self._process_information_per_pid = {}
    self._processes_per_pid = {}
    self._quiet_mode = False
    self._rpc_clients_per_pid = {}
    self._rpc_errors_per_pid = {}
    self._status_update_active = False
    self._status_update_thread = None
    self._storage_writer = None
    self._worker_memory_limit = definitions.DEFAULT_WORKER_MEMORY_LIMIT

  def _AbortJoin(self, timeout=None):
    """Aborts all registered processes by joining with the parent process.

    Args:
      timeout (int): number of seconds to wait for processes to join, where
          None represents no timeout.
    """
    for pid, process in self._processes_per_pid.items():
      logger.debug(f'Waiting for process: {process.name:s} (PID: {pid:d}).')
      process.join(timeout=timeout)
      if not process.is_alive():
        logger.debug(f'Process {process.name:s} (PID: {pid:d}) stopped.')

  def _AbortKill(self):
    """Aborts all registered processes by sending a SIGKILL or equivalent."""
    for pid, process in self._processes_per_pid.items():
      if not process.is_alive():
        continue

      logger.warning(f'Killing process: {process.name:s} (PID: {pid:d}).')
      self._KillProcess(pid)

  def _AbortTerminate(self):
    """Aborts all registered processes by sending a SIGTERM or equivalent."""
    for pid, process in self._processes_per_pid.items():
      if not process.is_alive():
        continue

      logger.warning(f'Terminating process: {process.name:s} (PID: {pid:d}).')
      process.terminate()

  def _CheckStatusWorkerProcess(self, pid):
    """Checks the status of a worker process.

    If a worker process is not responding the process is terminated and
    a replacement process is started.

    Args:
      pid (int): process ID (PID) of a registered worker process.

    Raises:
      KeyError: if the process is not registered with the engine.
    """
    # TODO: Refactor this method, simplify and separate concerns (monitoring
    # vs management).
    self._RaiseIfNotRegistered(pid)

    process = self._processes_per_pid[pid]

    process_status = self._QueryProcessStatus(process)
    if process_status is None:
      process_is_alive = False
    else:
      process_is_alive = True

    process_information = self._process_information_per_pid[pid]
    used_memory = process_information.GetUsedMemory() or 0

    if self._worker_memory_limit and used_memory > self._worker_memory_limit:
      logger.warning((
          f'Process: {process.name:s} (PID: {pid:d}) killed because it '
          f'exceeded the memory limit: {self._worker_memory_limit:d}.'))
      self._KillProcess(pid)

    if isinstance(process_status, dict):
      self._rpc_errors_per_pid[pid] = 0
      status_indicator = process_status.get('processing_status')

    else:
      rpc_errors = self._rpc_errors_per_pid.get(pid, 0) + 1
      self._rpc_errors_per_pid[pid] = rpc_errors

      if rpc_errors > self._MAXIMUM_RPC_ERRORS:
        process_is_alive = False

      if process_is_alive:
        rpc_port = process.rpc_port.value
        logger.warning((
            f'Unable to retrieve process: {process.name:s} (PID: {pid:d}) '
            f'status via RPC socket: http://localhost:{rpc_port:d}'))

        processing_status_string = 'RPC error'
        status_indicator = definitions.STATUS_INDICATOR_RUNNING
      else:
        processing_status_string = 'killed'
        status_indicator = definitions.STATUS_INDICATOR_KILLED

      process_status = {
          'processing_status': processing_status_string}

    self._UpdateProcessingStatus(pid, process_status, used_memory)

    # _UpdateProcessingStatus can also change the status of the worker,
    # So refresh the status if applicable.
    for worker_status in self._processing_status.workers_status:
      if worker_status.pid == pid:
        status_indicator = worker_status.status
        break

    if status_indicator in definitions.ERROR_STATUS_INDICATORS:
      logger.error((
          f'Process {process.name:s} (PID: {pid:d}) is not functioning '
          f'correctly. Status code: {status_indicator!s}.'))

      self._TerminateProcessByPid(pid)

      replacement_process = None
      replacement_process_name = f'Worker_{self._last_worker_number:02d}'
      for replacement_process_attempt in range(
          self._MAXIMUM_REPLACEMENT_RETRIES):
        attempt = replacement_process_attempt + 1
        logger.info((
            f'Attempt: {attempt:d} to start replacement worker process for '
            f'{process.name:s}'))

        replacement_process = self._StartWorkerProcess(replacement_process_name)
        if replacement_process:
          break

        time.sleep(self._REPLACEMENT_WORKER_RETRY_DELAY)

      if not replacement_process:
        logger.error((
            f'Unable to create replacement worker process for: '
            f'{process.name:s}'))

  def _KillProcess(self, pid):
    """Issues a SIGKILL or equivalent to the process.

    Args:
      pid (int): process identifier (PID).
    """
    if sys.platform.startswith('win'):
      process_terminate = 1
      handle = ctypes.windll.kernel32.OpenProcess(
          process_terminate, False, pid)
      ctypes.windll.kernel32.TerminateProcess(handle, -1)
      ctypes.windll.kernel32.CloseHandle(handle)

    else:
      try:
        os.kill(pid, signal.SIGKILL)
      except OSError as exception:
        logger.error(
            f'Unable to kill process {pid:d} with error: {exception!s}')

  def _QueryProcessStatus(self, process):
    """Queries a process to determine its status.

    Args:
      process (MultiProcessBaseProcess): process to query for its status.

    Returns:
      dict[str, str]: status values received from the worker process.
    """
    process_is_alive = process.is_alive()
    if not process_is_alive:
      return None

    rpc_client = self._rpc_clients_per_pid.get(process.pid)
    return rpc_client.CallFunction()

  def _RaiseIfNotMonitored(self, pid):
    """Raises if the process is not monitored by the engine.

    Args:
      pid (int): process identifier (PID).

    Raises:
      KeyError: if the process is not monitored by the engine.
    """
    if pid not in self._process_information_per_pid:
      raise KeyError(f'Process (PID: {pid:d}) not monitored by engine.')

  def _RaiseIfNotRegistered(self, pid):
    """Raises if the process is not registered with the engine.

    Args:
      pid (int): process identifier (PID).

    Raises:
      KeyError: if the process is not registered with the engine.
    """
    if pid not in self._processes_per_pid:
      raise KeyError(f'Process (PID: {pid:d}) not registered with engine.')

  def _RegisterProcess(self, process):
    """Registers a process with the engine.

    Args:
      process (MultiProcessBaseProcess): process.

    Raises:
      KeyError: if the process is already registered with the engine.
      ValueError: if the process is missing.
    """
    if process is None:
      raise ValueError('Missing process.')

    if process.pid in self._processes_per_pid:
      raise KeyError(
          f'Already managing process: {process.name!s} (PID: {process.pid:d})')

    self._processes_per_pid[process.pid] = process

  # pylint: disable=redundant-returns-doc
  @abc.abstractmethod
  def _StartWorkerProcess(self, process_name):
    """Creates, starts, monitors and registers a worker process.

    Args:
      process_name (str): process name.

    Returns:
      MultiProcessWorkerProcess: extraction worker process.
    """

  def _StartMonitoringProcess(self, process):
    """Starts monitoring a process.

    Args:
      process (MultiProcessBaseProcess): process.

    Raises:
      KeyError: if the process is not registered with the engine or
          if the process is already being monitored.
      OSError: if the RPC client cannot connect to the server.
      ValueError: if the process is missing.
    """
    if process is None:
      raise ValueError('Missing process.')

    pid = process.pid

    if pid in self._process_information_per_pid:
      raise KeyError(f'Already monitoring process (PID: {pid:d}).')

    if pid in self._rpc_clients_per_pid:
      raise KeyError(f'RPC client (PID: {pid:d}) already exists.')

    rpc_client = plaso_xmlrpc.XMLProcessStatusRPCClient()

    # Make sure that a worker process has started its RPC server.
    # The RPC port will be 0 if no server is available.
    rpc_port = process.rpc_port.value
    time_waited_for_process = 0.0
    while not rpc_port:
      time.sleep(0.1)
      rpc_port = process.rpc_port.value
      time_waited_for_process += 0.1

      if time_waited_for_process >= self._RPC_SERVER_TIMEOUT:
        raise OSError(
            f'RPC client unable to determine server (PID: {pid:d}) port.')

    hostname = 'localhost'

    if not rpc_client.Open(hostname, rpc_port):
      raise OSError((
          f'RPC client unable to connect to server (PID: {pid:d}) '
          f'http://{hostname:s}:{rpc_port:d}'))

    self._rpc_clients_per_pid[pid] = rpc_client
    self._process_information_per_pid[pid] = process_info.ProcessInfo(pid)

  def _StartStatusUpdateThread(self):
    """Starts the status update thread."""
    self._status_update_active = True
    self._status_update_thread = threading.Thread(
        name='Status update', target=self._StatusUpdateThreadMain)
    self._status_update_thread.start()

  def _StatusUpdateThreadMain(self):
    """Main function of the status update thread."""
    while self._status_update_active:
      self._UpdateStatus()
      time.sleep(self._status_update_interval)

  def _StopMonitoringProcess(self, process):
    """Stops monitoring a process.

    Args:
      process (MultiProcessBaseProcess): process.

    Raises:
      KeyError: if the process is not monitored.
      ValueError: if the process is missing.
    """
    if process is None:
      raise ValueError('Missing process.')

    pid = process.pid

    self._RaiseIfNotMonitored(pid)

    del self._process_information_per_pid[pid]

    rpc_client = self._rpc_clients_per_pid.get(pid)
    if rpc_client:
      rpc_client.Close()
      del self._rpc_clients_per_pid[pid]

    if pid in self._rpc_errors_per_pid:
      del self._rpc_errors_per_pid[pid]

    logger.debug('Stopped monitoring process: {process.name:s} (PID: {pid:d}).')

  def _StopMonitoringProcesses(self):
    """Stops monitoring all processes."""
    # We need to make a copy of the list of pids since we are changing
    # the dict in the loop.
    for pid in list(self._process_information_per_pid.keys()):
      self._RaiseIfNotRegistered(pid)
      process = self._processes_per_pid[pid]

      self._StopMonitoringProcess(process)

  def _StopStatusUpdateThread(self):
    """Stops the status update thread."""
    if self._status_update_thread:
      self._status_update_active = False
      if self._status_update_thread.is_alive():
        self._status_update_thread.join()
      self._status_update_thread = None

    # Update the status view one last time so we have the latest worker process
    # status information.
    self._UpdateStatus()

  def _TerminateProcessByPid(self, pid):
    """Terminate a process that's monitored by the engine.

    Args:
      pid (int): process identifier (PID).

    Raises:
      KeyError: if the process is not registered with and monitored by the
          engine.
    """
    self._RaiseIfNotRegistered(pid)

    process = self._processes_per_pid[pid]

    self._TerminateProcess(process)
    self._StopMonitoringProcess(process)

  def _TerminateProcess(self, process):
    """Terminate a process.

    Args:
      process (MultiProcessBaseProcess): process to terminate.
    """
    pid = process.pid
    logger.warning('Terminating process: (PID: {pid:d}).')
    process.terminate()

    # Wait for the process to exit.
    process.join(timeout=self._PROCESS_JOIN_TIMEOUT)

    if process.is_alive():
      logger.warning('Killing process: (PID: {pid:d}).')
      self._KillProcess(pid)

  @abc.abstractmethod
  def _UpdateProcessingStatus(self, pid, process_status, used_memory):
    """Updates the processing status.

    Args:
      pid (int): process identifier (PID) of the worker process.
      process_status (dict[str, object]): status values received from
          the worker process.
      used_memory (int): size of used memory in bytes.

    Raises:
      KeyError: if the process is not registered with the engine.
    """

  @abc.abstractmethod
  def _UpdateStatus(self):
    """Updates the status."""
