import inspect
import os
import random
import signal
import string
import sys
import traceback
from types import FunctionType

from async_generator._impl import ANextIter

from trio import Queue, WouldBlock, BrokenStreamError
from trio._highlevel_serve_listeners import _run_handler
from ._version import __version__
from trio.abc import Instrument
from trio.hazmat import current_task, Task


# inspiration: https://github.com/python-trio/trio/blob/master/notes-to-self/print-task-tree.py
# additional credit: python-trio/trio#429 (You are being ratelimited)


def walk_coro_stack(coro):
    while coro is not None:
        if hasattr(coro, "cr_frame"):
            # A real coroutine
            yield coro.cr_frame, coro.cr_frame.f_lineno
            coro = coro.cr_await
        elif isinstance(coro, ANextIter):
            # black hole
            return
        else:
            # A generator decorated with @types.coroutine
            yield coro.gi_frame, coro.gi_frame.f_lineno
            coro = coro.gi_yieldfrom


class Monitor(Instrument):
    """Represents a monitor; a simple way of monitoring the health of your
    Trio application using the Trio instrumentation API.

    The monitor protocol is a simple text-based protocol, accessible from a
    telnet client, for example.
    """
    MID_PREFIX = "|─ "
    MID_CONTINUE = "|  "
    END_PREFIX = "\─ "
    END_CONTINUE = " " * len(END_PREFIX)

    def __init__(self):
        self.authenticated = False

        # authentication for running code in a frame
        rand = random.SystemRandom()
        self.auth_pin = ''.join(
            rand.choice(string.ascii_letters) for x in range(0, 8)
        )

        self._is_monitoring = False
        # semi-arbitrary size, because otherwise we'll be dropping events
        # no clue how to make this better, alas.
        self._monitoring_queue = Queue(capacity=100)

    @staticmethod
    def get_root_task() -> Task:
        """Gets the current root task."""
        task = current_task()
        while task.parent_nursery is not None:
            task = task.parent_nursery.parent_task
        return task

    @staticmethod
    def flatten_tasks():
        """Gets a list of all tasks."""
        root = Monitor.get_root_task()
        tasks = [root]
        for child in root.child_nurseries:
            tasks.extend(Monitor.recursively_get_tasks(child))

        return tasks

    @staticmethod
    def recursively_get_tasks(nursery):
        """Recursively gets all tasks from a nursery."""
        tasks = []
        for task in nursery.child_tasks:
            tasks.append(task)
            for nursery in task.child_nurseries:
                tasks.extend(Monitor.recursively_get_tasks(nursery))

        return tasks

    @staticmethod
    def get_task_by_id(taskid: int):
        """Gets a task by ID."""
        for task in Monitor.flatten_tasks():
            if id(task) == taskid:
                return task

    # specific overrides
    def before_io_wait(self, timeout):
        if timeout == 0:
            return

        self._add_to_monitoring_queue(("before_io_wait", timeout))

    def after_io_wait(self, timeout):
        if timeout == 0:
            return

        self._add_to_monitoring_queue(("after_io_wait", timeout))

    def _add_to_monitoring_queue(self, item):
        if not self._is_monitoring:
            return

        if 'task' in item[0]:
            task = item[1]
            # idk how to make this better.
            if task.coro.cr_code == _run_handler.__code__:
                if task.coro.cr_frame is not None:
                    loc = task.coro.cr_frame.f_locals

                    # if it's our own handler, skip it!
                    if loc['handler'] == self.listen_on_stream:
                        return

        try:
            self._monitoring_queue.put_nowait(item)
        except WouldBlock:
            return

    async def listen_on_stream(self, stream):
        """Makes the monitor server listen on a stream.
        Use this as a callback from a listener.
        """
        return await self.main_loop(stream)

    async def main_loop(self, stream):
        """Runs the main loop of the monitor.
        """
        # send the banner
        version = __version__
        await stream.send_all(
            b"Connected to the Trio monitor, using "
            b"trio " + version.encode(encoding="ascii") + b"\n"
        )

        while True:
            await stream.send_all(b"trio> ")
            command = await stream.receive_some(2048)
            if command == b"":
                return

            command = command.decode("ascii").rstrip("\n").rstrip("\r")
            name, *args = command.split(" ")

            # special handling for closing
            if name in ["exit", "ex", "quit", "q", ":q"]:
                return await stream.aclose()

            # special handling for monitor
            if name in ["monitor", "mon", "m", "feed"]:
                try:
                    self._is_monitoring = True
                    return await self.do_monitor(stream)
                finally:
                    self._is_monitoring = False
                    # empty out the queue
                    self._monitoring_queue = Queue(capacity=100)

            try:
                fn = getattr(self, "command_{}".format(name))
            except AttributeError:
                await stream.send_all(
                    b"No such command: " + name.encode() + b"\n"
                )
                continue

            try:
                lines = await fn(*args)
            except Exception as e:
                messages = ["takes at most", "required positional argument"]

                if isinstance(e, TypeError) and \
                        any(x in ' '.join(e.args) for x in messages):
                    # hacky, but idk what else to do
                    await stream.send_all(
                        ' '.join(e.args).encode("ascii") + b"\n"
                    )
                    continue

                errormessage = type(e).__name__ + ": " + ' '.join(e.args)
                await stream.send_all(b"Error: " + errormessage.encode())
                raise

            await stream.send_all(
                "\n".join(lines).encode(encoding="ascii") + b"\n"
            )

    # monitor feed
    async def do_monitor(self, stream):
        """Livefeeds information about the running program."""
        prefix = "[FEED] "
        async for item in self._monitoring_queue:
            key = item[0]

            if key == "task_spawned":
                task = item[1]
                message = "Task spawned: {} ({})".format(task.name, id(task))

            elif key == "task_scheduled":
                task = item[1]
                message = "Task scheduled: {} ({})".format(task.name, id(task))
            elif key == "task_exited":
                task = item[1]
                message = "Task exited: {} ({})".format(task.name, id(task))

            elif key == "before_io_wait":
                timeout = item[1]
                message = "Waiting for IO (timeout: {:.3f})".format(timeout)

            elif key == "after_io_wait":
                timeout = item[1]
                message = "Done waiting for IO (timeout: {:.3f})" \
                    .format(timeout)

            elif key == "before_task_step":
                task = item[1]
                message = "Task stepping: {} ({})".format(task.name, id(task))

            elif key == "after_task_step":
                task = item[1]
                message = "Task finished stepping: {} ({})".format(
                    task.name, id(task)
                )

            else:
                message = "Unknown event: {}".format(key)

            message = prefix + message

            try:
                await stream.send_all(message.encode("ascii") + b'\n')
            except BrokenStreamError:  # client disconnected on us
                return

    # command definitions
    async def command_help(self):
        """Sends help."""
        name_rpad = 12

        def pred(i):
            return hasattr(i, "__name__") \
                   and i.__name__.startswith("command_")

        commands = inspect.getmembers(self, predicate=pred)
        lines = ["Commands:"]
        for name, command in commands:
            doc = inspect.getdoc(command).splitlines(keepends=False)[0]
            name = name.split("_", 1)[1]
            lines.append(name.ljust(name_rpad) + doc)

        return lines

    async def command_signal(self, signame: str):
        """Sends a signal to the server process."""
        signame = signame.upper()
        if not signame.startswith("SIG"):
            signame = "SIG{}".format(signame)

        try:
            tosend = getattr(signal, signame)
        except AttributeError:
            return ["Invalid signal: {}".format(signame)]

        os.kill(os.getpid(), tosend)
        return ["Signal sent successfully"]

    async def command_ps(self):
        """Gets the current list of tasks."""
        lines = []
        headers = ('ID', 'Name', 'Parent')
        widths = (15, 50, 15)
        header_line = []

        for name, width in zip(headers, widths):
            header_line.append(name.ljust(width))

        lines.append(' '.join(header_line))
        lines.append("-" * sum(widths))

        for task in self.flatten_tasks():
            if len(task.name) >= 50:
                name = task.name[:46] + "..."
            else:
                name = task.name

            if task.parent_nursery is None:
                parent = "N/A"
            else:
                parent = str(id(task.parent_nursery.parent_task))

            result = [str(id(task)).ljust(widths[0]), name.ljust(49), parent]

            lines.append(' '.join(result))

        return lines

    async def command_where(self, taskid):
        """Shows the stack frames of a task."""
        try:
            taskid = int(taskid)
        except ValueError:
            return ["Invalid task ID"]

        task = self.get_task_by_id(taskid)
        if task is None:
            return ["Invalid task ID"]

        summary = traceback.StackSummary.extract(walk_coro_stack(task.coro))
        lines = summary.format()
        lines = [line.rstrip('\n') for line in lines]
        return lines

    # stub commands
    async def command_monitor(self, *args):
        """Starts a live monitor feed."""
        return ["You shouldn't see this"]

    async def command_exit(self, *args):
        """Exits the monitor."""
        return ["You shouldn't see this"]


def _patch_monitor():
    def pred(i):
        return isinstance(i, FunctionType) and not i.__name__.startswith("__")

    for name, _ in inspect.getmembers(Instrument, predicate=pred):
        # 99% sure this is needed to bind the right name
        # otherwise it always uses the last name
        def bind(fname):
            def magic(self, *args):
                return self._add_to_monitoring_queue((fname, *args))

            return magic

        if getattr(Monitor, name) == getattr(Instrument, name):
            setattr(Monitor, name, bind(name))


_patch_monitor()
del _patch_monitor


def main():
    import argparse
    import telnetlib

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-a",
        "--address",
        default="127.0.0.1",
        help="The address to connect to"
    )
    parser.add_argument(
        "-p", "--port", default=14761, help="The port to connect to"
    )

    args = parser.parse_args()
    # TODO: Potentially wrap sys.stdin for better readline
    client = telnetlib.Telnet(host=args.address, port=args.port)
    client.interact()

    return 0


if __name__ == "__main__":
    sys.exit(main())
