"""Implements InterpreterPoolExecutor."""

import contextlib
import pickle
import textwrap
from . import thread as _thread
import _interpreters
import _interpqueues


class ExecutionFailed(_interpreters.InterpreterError):
    """An unhandled exception happened during execution."""

    def __init__(self, excinfo):
        msg = excinfo.formatted
        if not msg:
            if excinfo.type and excinfo.msg:
                msg = f'{excinfo.type.__name__}: {excinfo.msg}'
            else:
                msg = excinfo.type.__name__ or excinfo.msg
        super().__init__(msg)
        self.excinfo = excinfo

    def __str__(self):
        try:
            formatted = self.excinfo.errdisplay
        except Exception:
            return super().__str__()
        else:
            return textwrap.dedent(f"""
{super().__str__()}

Uncaught in the interpreter:

{formatted}
                """.strip())


UNBOUND = 2  # error; this should not happen.


class WorkerContext(_thread.WorkerContext):

    @classmethod
    def prepare(cls, initializer, initargs, shared):
        def resolve_task(fn, args, kwargs):
            if isinstance(fn, str):
                # XXX Circle back to this later.
                raise TypeError('scripts not supported')
                if args or kwargs:
                    raise ValueError(f'a script does not take args or kwargs, got {args!r} and {kwargs!r}')
                data = textwrap.dedent(fn)
                kind = 'script'
                # Make sure the script compiles.
                # Ideally we wouldn't throw away the resulting code
                # object.  However, there isn't much to be done until
                # code objects are shareable and/or we do a better job
                # of supporting code objects in _interpreters.exec().
                compile(data, '<string>', 'exec')
            else:
                # Functions defined in the __main__ module can't be pickled,
                # so they can't be used here.  In the future, we could possibly
                # borrow from multiprocessing to work around this.
                data = pickle.dumps((fn, args, kwargs))
                kind = 'function'
            return (data, kind)

        if initializer is not None:
            try:
                initdata = resolve_task(initializer, initargs, {})
            except ValueError:
                if isinstance(initializer, str) and initargs:
                    raise ValueError(f'an initializer script does not take args, got {initargs!r}')
                raise  # re-raise
        else:
            initdata = None
        def create_context():
            return cls(initdata, shared)
        return create_context, resolve_task

    @classmethod
    @contextlib.contextmanager
    def _capture_exc(cls, resultsid):
        try:
            yield
        except BaseException as exc:
            # Send the captured exception out on the results queue,
            # but still leave it unhandled for the interpreter to handle.
            err = pickle.dumps(exc)
            _interpqueues.put(resultsid, (None, err), 1, UNBOUND)
            raise  # re-raise

    @classmethod
    def _send_script_result(cls, resultsid):
        _interpqueues.put(resultsid, (None, None), 0, UNBOUND)

    @classmethod
    def _call(cls, func, args, kwargs, resultsid):
        with cls._capture_exc(resultsid):
            res = func(*args or (), **kwargs or {})
        # Send the result back.
        try:
            _interpqueues.put(resultsid, (res, None), 0, UNBOUND)
        except _interpreters.NotShareableError:
            res = pickle.dumps(res)
            _interpqueues.put(resultsid, (res, None), 1, UNBOUND)

    @classmethod
    def _call_pickled(cls, pickled, resultsid):
        with cls._capture_exc(resultsid):
            fn, args, kwargs = pickle.loads(pickled)
        cls._call(fn, args, kwargs, resultsid)

    def __init__(self, initdata, shared=None):
        self.initdata = initdata
        self.shared = dict(shared) if shared else None
        self.interpid = None
        self.resultsid = None

    def __del__(self):
        if self.interpid is not None:
            self.finalize()

    def _exec(self, script):
        assert self.interpid is not None
        excinfo = _interpreters.exec(self.interpid, script, restrict=True)
        if excinfo is not None:
            raise ExecutionFailed(excinfo)

    def initialize(self):
        assert self.interpid is None, self.interpid
        self.interpid = _interpreters.create(reqrefs=True)
        try:
            _interpreters.incref(self.interpid)

            maxsize = 0
            fmt = 0
            self.resultsid = _interpqueues.create(maxsize, fmt, UNBOUND)

            self._exec(f'from {__name__} import WorkerContext')

            if self.shared:
                _interpreters.set___main___attrs(
                                    self.interpid, self.shared, restrict=True)

            if self.initdata:
                self.run(self.initdata)
        except BaseException:
            self.finalize()
            raise  # re-raise

    def finalize(self):
        interpid = self.interpid
        resultsid = self.resultsid
        self.resultsid = None
        self.interpid = None
        if resultsid is not None:
            try:
                _interpqueues.destroy(resultsid)
            except _interpqueues.QueueNotFoundError:
                pass
        if interpid is not None:
            try:
                _interpreters.decref(interpid)
            except _interpreters.InterpreterNotFoundError:
                pass

    def run(self, task):
        data, kind = task
        if kind == 'script':
            raise NotImplementedError('script kind disabled')
            script = f"""
with WorkerContext._capture_exc({self.resultsid}):
{textwrap.indent(data, '    ')}
WorkerContext._send_script_result({self.resultsid})"""
        elif kind == 'function':
            script = f'WorkerContext._call_pickled({data!r}, {self.resultsid})'
        else:
            raise NotImplementedError(kind)

        try:
            self._exec(script)
        except ExecutionFailed as exc:
            exc_wrapper = exc
        else:
            exc_wrapper = None

        # Return the result, or raise the exception.
        while True:
            try:
                obj = _interpqueues.get(self.resultsid)
            except _interpqueues.QueueNotFoundError:
                raise  # re-raise
            except _interpqueues.QueueError:
                continue
            except ModuleNotFoundError:
                # interpreters.queues doesn't exist, which means
                # QueueEmpty doesn't.  Act as though it does.
                continue
            else:
                break
        (res, excdata), pickled, unboundop = obj
        assert unboundop is None, unboundop
        if excdata is not None:
            assert res is None, res
            assert pickled
            assert exc_wrapper is not None
            exc = pickle.loads(excdata)
            raise exc from exc_wrapper
        return pickle.loads(res) if pickled else res


class BrokenInterpreterPool(_thread.BrokenThreadPool):
    """
    Raised when a worker thread in an InterpreterPoolExecutor failed initializing.
    """


class InterpreterPoolExecutor(_thread.ThreadPoolExecutor):

    BROKEN = BrokenInterpreterPool

    @classmethod
    def prepare_context(cls, initializer, initargs, shared):
        return WorkerContext.prepare(initializer, initargs, shared)

    def __init__(self, max_workers=None, thread_name_prefix='',
                 initializer=None, initargs=(), shared=None):
        """Initializes a new InterpreterPoolExecutor instance.

        Args:
            max_workers: The maximum number of interpreters that can be used to
                execute the given calls.
            thread_name_prefix: An optional name prefix to give our threads.
            initializer: A callable or script used to initialize
                each worker interpreter.
            initargs: A tuple of arguments to pass to the initializer.
            shared: A mapping of shareabled objects to be inserted into
                each worker interpreter.
        """
        super().__init__(max_workers, thread_name_prefix,
                         initializer, initargs, shared=shared)