import multiprocessing as mp
import os
import pathlib
import sys
import threading
import traceback
import warnings

import cloudpickle
import psutil

from . import utils


class Context:

  def __init__(self):
    self.initfns = []
    self.resolver = None
    self.errfile = None
    self.interval = 20
    self.clientkw = {}
    self.serverkw = {}
    self.printlock = threading.Lock()
    self.done = threading.Event()
    self.watcher = None

  @property
  def mp(self):
    return mp.get_context('spawn')

  def options(self):
    return {
        'resolver': self.resolver and cloudpickle.dumps(self.resolver),
        'errfile': self.errfile,
        'interval': self.interval,
        'initfns': self.initfns,
        'clientkw': self.clientkw,
        'serverkw': self.serverkw,
    }

  def setup(
      self,
      resolver=None,
      errfile=None,
      interval=None,
      initfns=None,
      clientkw=None,
      serverkw=None,
      host=None,
      ipv6=None,
  ):

    if resolver:
      if isinstance(resolver, bytes):
        resolver = cloudpickle.loads(resolver)
      assert callable(resolver)
      self.resolver = resolver

    if errfile:
      if isinstance(errfile, str):
        errfile = pathlib.Path(errfile)
      assert hasattr(errfile, 'exists') and hasattr(errfile, 'write_text')
      self.errfile = errfile
      self._check_errfile()
      self._install_excepthook()

    if interval:
      assert isinstance(interval, (int, float))
      self.interval = interval

    if initfns:
      self.initfns = []
      for fn in initfns:
        self.initfn(fn, call_now=True)

    if clientkw is not None:
      assert isinstance(clientkw, dict)
      self.clientkw = clientkw

    if serverkw is not None:
      assert isinstance(serverkw, dict)
      self.serverkw = serverkw

    if host is not None:
      self.serverkw['host'] = host

    if ipv6 is not None:
      self.clientkw['ipv6'] = ipv6
      self.serverkw['ipv6'] = ipv6

    if self.errfile and not self.watcher:
      self.watcher = threading.Thread(target=self._watcher, daemon=True)
      self.watcher.start()

  def initfn(self, fn, call_now=True):
    if isinstance(fn, bytes):
      pkl, fn = fn, cloudpickle.loads(fn)
    else:
      pkl, fn = cloudpickle.dumps(fn), fn
    self.initfns.append(pkl)
    call_now and fn()

  def print(self, name, *args, color=None):
    assert args
    if color:
      style = utils.style(color=color)
      reset = utils.style(reset=True)
    else:
      style, reset = '', ''
    with self.printlock:
      print(style + f'[{name}]' + reset, *args)

  def error(self, e, name=None):
    typ, tb = type(e), e.__traceback__
    summary = list(traceback.format_exception_only(typ, e))[0].strip('\n')
    long = ''.join(traceback.format_exception(typ, e, tb)).strip('\n')
    message = f"Error in '{name}' ({summary}):\n{long}"
    with self.printlock:
      style = utils.style(color='red')
      reset = utils.style(reset=True)
      print(style + '\n---\n' + message + reset, file=sys.stderr)
    if self.errfile:
      self.errfile.write_text(message)
      print(f'Wrote errorfile: {self.errfile}', file=sys.stderr)

  def shutdown(self, exitcode):
    children = list(psutil.Process(os.getpid()).children(recursive=True))
    utils.kill_proc(children, timeout=1)
    # TODO
    # if exitcode == 0:
    #   for child in self.children(threading.main_thread()):
    #     child.kill()
    #   os._exit(0)
    # else:
    os._exit(exitcode)

  def close(self):
    self.done.set()
    if self.watcher:
      self.watcher.join()

  def add_worker(self, worker):
    assert hasattr(worker, 'kill')
    current = threading.current_thread()
    if current == threading.main_thread():
      return
    if hasattr(worker, 'thread'):
      assert current != worker.thread
    try:
      current.children.append(worker)
    except AttributeError:
      warnings.warn(
          'Using Portal from plain Python threads is discouraged because ' +
          'they can cause hangs during shutdown.')

  def children(self, thread):
    current = thread or threading.current_thread()
    if current == threading.main_thread():
      return []
    return current.children

  def _watcher(self):
    while not self.done.wait(self.interval):
      self._check_errfile()

  def _install_excepthook(self):
    existing = sys.excepthook
    def patched(typ, val, tb):
      if self.errfile:
        message = ''.join(traceback.format_exception(typ, val, tb)).strip('\n')
        self.errfile.write_text(message)
        print(f'Wrote errorfile: {self.errfile}', file=sys.stderr)
      return existing(typ, val, tb)
    sys.excepthook = patched

  def _check_errfile(self):
    if self.errfile and self.errfile.exists():
      message = f'Shutting down due to error file: {self.errfile}'
      print(message, file=sys.stderr)
      self.shutdown(exitcode=2)


context = Context()


def initfn(fn, call_now=True):
  context.initfn(fn, call_now)


def setup(**kwargs):
  context.setup(**kwargs)


def reset():
  global context
  context.close()
  context = Context()