import asyncio
import contextlib
import functools
import sys
import traceback

from codeop import CommandCompiler
from io import BytesIO, StringIO


__all__ = ['start_manhole']


class StatefulCommandCompiler(CommandCompiler):
    """A command compiler that buffers input until a full command is available."""

    def __init__(self):
        super().__init__()
        self.buf = BytesIO()

    def is_partial_command(self):
        return bool(self.buf.getvalue())

    def __call__(self, source, **kwargs):
        buf = self.buf
        if self.is_partial_command():
            buf.write(b'\n')
        buf.write(source)

        code = self.buf.getvalue().decode('utf8')

        codeobj = super().__call__(code, **kwargs)

        if codeobj:
            self.reset()
        return codeobj

    def reset(self):
        self.buf.seek(0)
        self.buf.truncate(0)


class InteractiveInterpreter:
    """An interactive asynchronous interpreter."""

    def __init__(self, namespace, banner, loop):
        self.namespace = namespace
        self.banner = self.get_banner(banner)
        self.compiler = StatefulCommandCompiler()
        self.loop = loop

    def get_banner(self, banner):
        if isinstance(banner, bytes):
            return banner
        elif isinstance(banner, str):
            return banner.encode('utf8')
        elif banner is None:
            return b''
        else:
            raise ValueError("Cannot handle unknown banner type {!}, expected str or bytes".format(banner.__class__.__name__))

    def attempt_compile(self, line):
        return self.compiler(line)

    @asyncio.coroutine
    def send_exception(self):
        """When an exception has occurred, write the traceback to the user."""
        self.compiler.reset()

        exc = traceback.format_exc()
        self.writer.write(exc.encode('utf8'))

        yield from self.writer.drain()

    @asyncio.coroutine
    def attempt_exec(self, codeobj, namespace):
        with contextlib.redirect_stdout(StringIO()) as buf:
            value = yield from self._real_exec(codeobj, namespace)

        return value, buf.getvalue()

    @asyncio.coroutine
    def _real_exec(self, codeobj, namespace):
        yield  # quick hack to pretend to be a coroutine, which attempt_exec expects
        return eval(codeobj, namespace)

    @asyncio.coroutine
    def handle_one_command(self):
        """Process a single command. May have many lines."""

        while True:
            yield from self.write_prompt()
            codeobj = yield from self.read_command()

            if codeobj is not None:
                yield from self.run_command(codeobj)

    @asyncio.coroutine
    def run_command(self, codeobj):
        """Execute a compiled code object, and write the output back to the client."""
        try:
            value, stdout = yield from self.attempt_exec(codeobj, self.namespace)
        except Exception:
            yield from self.send_exception()
            return
        else:
            yield from self.send_output(value, stdout)

    @asyncio.coroutine
    def write_prompt(self):
        writer = self.writer

        if self.compiler.is_partial_command():
            writer.write(sys.ps2.encode('utf8'))
        else:
            writer.write(sys.ps1.encode('utf8'))

        yield from writer.drain()

    @asyncio.coroutine
    def read_command(self):
        """Read a command from the user line by line.