import atexit
import errno
import traceback

import cloudpickle
import psutil

from . import contextlib
from . import utils


class Process:

  """
  This process implementation extends the standard Python process as follows:

  1. It provides a kill() method that not only terminates the process itself
  but also all its nested child processes. This prevents subprocesses from
  lingering in the background after the Python program has ended.

  2. The process terminates its nested child processes when it encounters an
  error. This prevents lingering subprocesses on error.

  3. When the parent process encounters an error, its subprocess will be killed
  via `atexit`, preventing hangs.

  4. It inherits the context object of its parent process, which provides
  cloudpickled initializer functions for each nested child process and error
  file watching for global shutdown on error.

  5. Standard Python subprocesses do not reliably return the running() state of
  the process. This class makes it more reliable.
  """

  def __init__(self, fn, *args, name=None, start=False):
    name = name or getattr(fn, '__name__', 'process')
    fn = cloudpickle.dumps(fn)
    options = contextlib.context.options()
    self.ready = contextlib.context.mp.Barrier(2)
    self.process = contextlib.context.mp.Process(
        target=self._wrapper, name=name,
        args=(options, self.ready, name, fn, args))
    self.started = False
    self.killed = False
    self.thepid = None
    atexit.register(self.kill)
    contextlib.context.add_worker(self)
    start and self.start()

  @property
  def name(self):
    return self.process.name

  @property
  def pid(self):
    return self.thepid

  @property
  def running(self):
    if not self.started:
      return False
    if not self.process.is_alive():
      return False
    return utils.proc_alive(self.pid)

  @property
  def exitcode(self):
    exitcode = self.process.exitcode
    if self.killed and exitcode is None:
      return -9
    return exitcode

  def start(self):
    assert not self.started
    self.started = True
    self.process.start()
    self.ready.wait()
    self.thepid = self.process.pid
    assert self.thepid is not None
    return self

  def join(self, timeout=None):
    assert self.started
    self.process.join(timeout)
    return self

  def kill(self, timeout=1):
    assert self.started
    try:
      children = list(psutil.Process(self.pid).children(recursive=True))
    except psutil.NoSuchProcess:
      children = []
    try:
      self.process.terminate()
      self.process.join(timeout)
      self.process.kill()
      utils.kill_proc(children, timeout)
    except OSError as e:
      if e.errno != errno.ESRCH:
        contextlib.context.error(e, self.name)
        contextlib.context.shutdown(exitcode=1)
    self.killed = True
    return self

  def __repr__(self):
    attrs = ('name', 'pid', 'running', 'exitcode')
    attrs = [f'{k}={getattr(self, k)}' for k in attrs]
    return 'Process(' + ', '.join(attrs) + ')'

  @staticmethod
  def _wrapper(options, ready, name, fn, args):
    exitcode = 0
    try:
      ready.wait()
      contextlib.setup(**options)
      fn = cloudpickle.loads(fn)
      exitcode = fn(*args)
      exitcode = exitcode if isinstance(exitcode, int) else 0
    except (SystemExit, KeyboardInterrupt) as e:
      compact = traceback.format_tb(e.__traceback__)
      compact = '\n'.join([line.split('\n', 1)[0] for line in compact])
      print(f"Killed process '{name}' at:\n{compact}")
      exitcode = 2
    except Exception as e:
      contextlib.context.error(e, name)
      exitcode = 1
    finally:
      contextlib.context.shutdown(exitcode)