"""
Base class for protocol layers.
"""

import collections
import textwrap
from abc import abstractmethod
from collections.abc import Callable
from collections.abc import Generator
from dataclasses import dataclass
from logging import DEBUG
from typing import Any
from typing import ClassVar
from typing import NamedTuple
from typing import TypeVar

from mitmproxy.connection import Connection
from mitmproxy.proxy import commands
from mitmproxy.proxy import events
from mitmproxy.proxy.commands import Command
from mitmproxy.proxy.commands import StartHook
from mitmproxy.proxy.context import Context

T = TypeVar("T")
CommandGenerator = Generator[Command, Any, T]
"""
A function annotated with CommandGenerator[bool] may yield commands and ultimately return a boolean value.
"""


MAX_LOG_STATEMENT_SIZE = 2048
"""Maximum size of individual log statements before they will be truncated."""


class Paused(NamedTuple):
    """
    State of a layer that's paused because it is waiting for a command reply.
    """

    command: commands.Command
    generator: CommandGenerator


class Layer:
    """
    The base class for all protocol layers.

    Layers interface with their child layer(s) by calling .handle_event(event),
    which returns a list (more precisely: a generator) of commands.
    Most layers do not implement .directly, but instead implement ._handle_event, which
    is called by the default implementation of .handle_event.
    The default implementation of .handle_event allows layers to emulate blocking code:
    When ._handle_event yields a command that has its blocking attribute set to True, .handle_event pauses
    the execution of ._handle_event and waits until it is called with the corresponding CommandCompleted event.
    All events encountered in the meantime are buffered and replayed after execution is resumed.

    The result is code that looks like blocking code, but is not blocking:

        def _handle_event(self, event):
            err = yield OpenConnection(server)  # execution continues here after a connection has been established.

    Technically this is very similar to how coroutines are implemented.
    """

    __last_debug_message: ClassVar[str] = ""
    context: Context
    _paused: Paused | None
    """
    If execution is currently paused, this attribute stores the paused coroutine
    and the command for which we are expecting a reply.
    """
    _paused_event_queue: collections.deque[events.Event]
    """
    All events that have occurred since execution was paused.
    These will be replayed to ._child_layer once we resume.
    """
    debug: str | None = None
    """
    Enable debug logging by assigning a prefix string for log messages.
    Different amounts of whitespace for different layers work well.
    """

    def __init__(self, context: Context) -> None:
        self.context = context
        self.context.layers.append(self)
        self._paused = None
        self._paused_event_queue = collections.deque()

        show_debug_output = getattr(context.options, "proxy_debug", False)
        if show_debug_output:  # pragma: no cover
            self.debug = "  " * len(context.layers)

    def __repr__(self):
        statefun = getattr(self, "state", self._handle_event)
        state = getattr(statefun, "__name__", "")
        state = state.replace("state_", "")
        if state == "_handle_event":
            state = ""
        else:
            state = f"state: {state}"
        return f"{type(self).__name__}({state})"

    def __debug(self, message):
        """yield a Log command indicating what message is passing through this layer."""
        if len(message) > MAX_LOG_STATEMENT_SIZE:
            message = message[:MAX_LOG_STATEMENT_SIZE] + "…"
        if Layer.__last_debug_message == message:
            message = message.split("\n", 1)[0].strip()
            if len(message) > 256:
                message = message[:256] + "…"
        else:
            Layer.__last_debug_message = message
        assert self.debug is not None
        return commands.Log(textwrap.indent(message, self.debug), DEBUG)

    @property
    def stack_pos(self) -> str:
        """repr() for this layer and all its parent layers, only useful for debugging."""
        try:
            idx = self.context.layers.index(self)
        except ValueError:
            return repr(self)
        else:
            return " >> ".join(repr(x) for x in self.context.layers[: idx + 1])

    @abstractmethod
    def _handle_event(self, event: events.Event) -> CommandGenerator[None]:
        """Handle a proxy server event"""
        yield from ()  # pragma: no cover

    def handle_event(self, event: events.Event) -> CommandGenerator[None]:
        if self._paused:
            # did we just receive the reply we were waiting for?
            pause_finished = (
                isinstance(event, events.CommandCompleted)
                and event.command is self._paused.command
            )
            if self.debug is not None:
                yield self.__debug(f"{'>>' if pause_finished else '>!'} {event}")
            if pause_finished:
                assert isinstance(event, events.CommandCompleted)
                yield from self.__continue(event)
            else:
                self._paused_event_queue.append(event)
        else:
            if self.debug is not None:
                yield self.__debug(f">> {event}")
            command_generator = self._handle_event(event)
            send = None

            # inlined copy of __process to reduce call stack.
            # <✂✂✂>
            try:
                # Run ._handle_event to the next yield statement.
                # If you are not familiar with generators and their .send() method,
                # https://stackoverflow.com/a/12638313/934719 has a good explanation.
                command = command_generator.send(send)
            except StopIteration:
                return

            while True:
                if self.debug is not None:
                    if not isinstance(command, commands.Log):
                        yield self.__debug(f"<< {command}")
                if command.blocking is True:
                    # We only want this layer to block, the outer layers should not block.
                    # For example, take an HTTP/2 connection: If we intercept one particular request,
                    # we don't want all other requests in the connection to be blocked a well.
                    # We signal to outer layers that this command is already handled by assigning our layer to
                    # `.blocking` here (upper layers explicitly check for `is True`).
                    command.blocking = self
                    self._paused = Paused(
                        command,
                        command_generator,
                    )
                    yield command
                    return
                else:
                    yield command
                    try:
                        command = next(command_generator)
                    except StopIteration:
                        return
            # </✂✂✂>

    def __process(self, command_generator: CommandGenerator, send=None):
        """
        Yield commands from a generator.
        If a command is blocking, execution is paused and this function returns without
        processing any further commands.
        """
        try:
            # Run ._handle_event to the next yield statement.
            # If you are not familiar with generators and their .send() method,
            # https://stackoverflow.com/a/12638313/934719 has a good explanation.
            command = command_generator.send(send)
        except StopIteration:
            return

        while True:
            if self.debug is not None:
                if not isinstance(command, commands.Log):
                    yield self.__debug(f"<< {command}")
            if command.blocking is True:
                # We only want this layer to block, the outer layers should not block.
                # For example, take an HTTP/2 connection: If we intercept one particular request,
                # we don't want all other requests in the connection to be blocked a well.
                # We signal to outer layers that this command is already handled by assigning our layer to
                # `.blocking` here (upper layers explicitly check for `is True`).
                command.blocking = self
                self._paused = Paused(
                    command,
                    command_generator,
                )
                yield command
                return
            else:
                yield command
                try:
                    command = next(command_generator)
                except StopIteration:
                    return

    def __continue(self, event: events.CommandCompleted):
        """
        Continue processing events after being paused.
        The tricky part here is that events in the event queue may trigger commands which again pause the execution,
        so we may not be able to process the entire queue.
        """
        assert self._paused is not None
        command_generator = self._paused.generator
        self._paused = None
        yield from self.__process(command_generator, event.reply)

        while not self._paused and self._paused_event_queue:
            ev = self._paused_event_queue.popleft()
            if self.debug is not None:
                yield self.__debug(f"!> {ev}")
            command_generator = self._handle_event(ev)
            yield from self.__process(command_generator)


mevents = (
    events  # alias here because autocomplete above should not have aliased version.
)


class NextLayer(Layer):
    layer: Layer | None
    """The next layer. To be set by an addon."""

    events: list[mevents.Event]
    """All events that happened before a decision was made."""

    _ask_on_start: bool

    def __init__(self, context: Context, ask_on_start: bool = False) -> None:
        super().__init__(context)
        self.context.layers.remove(self)
        self.layer = None
        self.events = []
        self._ask_on_start = ask_on_start
        self._handle: Callable[[mevents.Event], CommandGenerator[None]] | None = None

    def __repr__(self):
        return f"NextLayer:{self.layer!r}"

    def handle_event(self, event: mevents.Event):
        if self._handle is not None:
            yield from self._handle(event)
        else:
            yield from super().handle_event(event)

    def _handle_event(self, event: mevents.Event):
        self.events.append(event)

        # We receive new data. Let's find out if we can determine the next layer now?
        if self._ask_on_start and isinstance(event, events.Start):
            yield from self._ask()
        elif (
            isinstance(event, mevents.ConnectionClosed)
            and event.connection == self.context.client
        ):
            # If we have not determined the next protocol yet and the client already closes the connection,
            # we abort everything.
            yield commands.CloseConnection(self.context.client)
        elif isinstance(event, mevents.DataReceived):
            # For now, we only ask if we have received new data to reduce hook noise.
            yield from self._ask()

    def _ask(self):
        """
        Manually trigger a next_layer hook.
        The only use at the moment is to make sure that the top layer is initialized.
        """
        yield NextLayerHook(self)

        # Has an addon decided on the next layer yet?
        if self.layer:
            if self.debug:
                yield commands.Log(f"{self.debug}[nextlayer] {self.layer!r}", DEBUG)
            for e in self.events:
                yield from self.layer.handle_event(e)
            self.events.clear()

            # Why do we need three assignments here?
            #  1. When this function here is invoked we may have paused events. Those should be
            #     forwarded to the sublayer right away, so we reassign ._handle_event.
            #  2. This layer is not needed anymore, so we directly reassign .handle_event.
            #  3. Some layers may however still have a reference to the old .handle_event.
            #     ._handle is just an optimization to reduce the callstack in these cases.
            self.handle_event = self.layer.handle_event  # type: ignore
            self._handle_event = self.layer.handle_event  # type: ignore
            self._handle = self.layer.handle_event

    # Utility methods for whoever decides what the next layer is going to be.
    def data_client(self):
        return self._data(self.context.client)

    def data_server(self):
        return self._data(self.context.server)

    def _data(self, connection: Connection):
        data = (
            e.data
            for e in self.events
            if isinstance(e, mevents.DataReceived) and e.connection == connection
        )
        return b"".join(data)


@dataclass
class NextLayerHook(StartHook):
    """
    Network layers are being switched. You may change which layer will be used by setting data.layer.

    (by default, this is done by mitmproxy.addons.NextLayer)
    """

    data: NextLayer