import time
from enum import Enum, auto
from typing import Optional, Union

from mitmproxy import connection
from mitmproxy.proxy import commands, context, events, layer
from mitmproxy.proxy.layer import Layer


class TunnelState(Enum):
    INACTIVE = auto()
    ESTABLISHING = auto()
    OPEN = auto()
    CLOSED = auto()


class TunnelLayer(layer.Layer):
    """
    A specialized layer that simplifies the implementation of tunneling protocols such as SOCKS, upstream HTTP proxies,
    or TLS.
    """

    child_layer: layer.Layer
    tunnel_connection: connection.Connection
    """The 'outer' connection which provides the tunnel protocol I/O"""
    conn: connection.Connection
    """The 'inner' connection which provides data I/O"""
    tunnel_state: TunnelState = TunnelState.INACTIVE
    command_to_reply_to: Optional[commands.OpenConnection] = None
    _event_queue: list[events.Event]
    """
    If the connection already exists when we receive the start event,
    we buffer commands until we have established the tunnel.
    """

    def __init__(
        self,
        context: context.Context,
        tunnel_connection: connection.Connection,
        conn: connection.Connection,
    ):
        super().__init__(context)
        self.tunnel_connection = tunnel_connection
        self.conn = conn
        self.child_layer = layer.NextLayer(self.context)
        self._event_queue = []

    def __repr__(self):
        return f"{type(self).__name__}({self.tunnel_state.name.lower()})"

    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:
        if isinstance(event, events.Start):
            if self.tunnel_connection.state is not connection.ConnectionState.CLOSED:
                # we might be in the interesting state here where the connection is already half-closed,
                # for example because next_layer buffered events and the client disconnected in the meantime.
                # we still expect a close event to arrive, so we carry on here as normal for now.
                self.tunnel_state = TunnelState.ESTABLISHING
                yield from self.start_handshake()
            yield from self.event_to_child(event)
        elif (
            isinstance(event, events.ConnectionEvent)
            and event.connection == self.tunnel_connection
        ):
            if isinstance(event, events.DataReceived):
                if self.tunnel_state is TunnelState.ESTABLISHING:
                    done, err = yield from self.receive_handshake_data(event.data)
                    if done:
                        if self.conn != self.tunnel_connection:
                            self.conn.state = connection.ConnectionState.OPEN
                            self.conn.timestamp_start = time.time()
                    if err:
                        if self.conn != self.tunnel_connection:
                            self.conn.state = connection.ConnectionState.CLOSED
                            self.conn.timestamp_start = time.time()
                        yield from self.on_handshake_error(err)
                    if done or err:
                        yield from self._handshake_finished(err)
                else:
                    yield from self.receive_data(event.data)
            elif isinstance(event, events.ConnectionClosed):
                if self.conn != self.tunnel_connection:
                    self.conn.state &= ~connection.ConnectionState.CAN_READ
                    self.conn.timestamp_end = time.time()
                if self.tunnel_state is TunnelState.OPEN:
                    yield from self.receive_close()
                elif self.tunnel_state is TunnelState.ESTABLISHING:
                    err = "connection closed"
                    yield from self.on_handshake_error(err)
                    yield from self._handshake_finished(err)
                self.tunnel_state = TunnelState.CLOSED
            else:  # pragma: no cover
                raise AssertionError(f"Unexpected event: {event}")
        else:
            yield from self.event_to_child(event)

    def _handshake_finished(self, err: Optional[str]):
        if err:
            self.tunnel_state = TunnelState.CLOSED
        else:
            self.tunnel_state = TunnelState.OPEN
        if self.command_to_reply_to:
            yield from self.event_to_child(
                events.OpenConnectionCompleted(self.command_to_reply_to, err)
            )
            self.command_to_reply_to = None
        else:
            for evt in self._event_queue:
                yield from self.event_to_child(evt)
            self._event_queue.clear()

    def event_to_child(self, event: events.Event) -> layer.CommandGenerator[None]:
        if (
            self.tunnel_state is TunnelState.ESTABLISHING
            and not self.command_to_reply_to
        ):
            self._event_queue.append(event)
            return
        for command in self.child_layer.handle_event(event):
            if (
                isinstance(command, commands.ConnectionCommand)
                and command.connection == self.conn
            ):
                if isinstance(command, commands.SendData):
                    yield from self.send_data(command.data)
                elif isinstance(command, commands.CloseConnection):
                    if self.conn != self.tunnel_connection:
                        if command.half_close:
                            self.conn.state &= ~connection.ConnectionState.CAN_WRITE
                        else:
                            self.conn.state = connection.ConnectionState.CLOSED
                    yield from self.send_close(command.half_close)
                elif isinstance(command, commands.OpenConnection):
                    # create our own OpenConnection command object that blocks here.
                    self.command_to_reply_to = command
                    self.tunnel_state = TunnelState.ESTABLISHING
                    err = yield commands.OpenConnection(self.tunnel_connection)
                    if err:
                        yield from self.event_to_child(
                            events.OpenConnectionCompleted(command, err)
                        )
                        self.tunnel_state = TunnelState.CLOSED
                    else:
                        yield from self.start_handshake()
                else:  # pragma: no cover
                    raise AssertionError(f"Unexpected command: {command}")
            else:
                yield command

    def start_handshake(self) -> layer.CommandGenerator[None]:
        yield from self._handle_event(events.DataReceived(self.tunnel_connection, b""))

    def receive_handshake_data(
        self, data: bytes
    ) -> layer.CommandGenerator[tuple[bool, Optional[str]]]:
        """returns a (done, err) tuple"""
        yield from ()
        return True, None

    def on_handshake_error(self, err: str) -> layer.CommandGenerator[None]:
        """Called if either receive_handshake_data returns an error or we receive a close during handshake."""
        yield commands.CloseConnection(self.tunnel_connection)

    def receive_data(self, data: bytes) -> layer.CommandGenerator[None]:
        yield from self.event_to_child(events.DataReceived(self.conn, data))

    def receive_close(self) -> layer.CommandGenerator[None]:
        yield from self.event_to_child(events.ConnectionClosed(self.conn))

    def send_data(self, data: bytes) -> layer.CommandGenerator[None]:
        yield commands.SendData(self.tunnel_connection, data)

    def send_close(self, half_close: bool) -> layer.CommandGenerator[None]:
        yield commands.CloseConnection(self.tunnel_connection, half_close=half_close)


class LayerStack:
    def __init__(self) -> None:
        self._stack: list[Layer] = []

    def __getitem__(self, item: int) -> Layer:
        return self._stack.__getitem__(item)

    def __truediv__(self, other: Union[Layer, "LayerStack"]) -> "LayerStack":
        if isinstance(other, Layer):
            if self._stack:
                self._stack[-1].child_layer = other  # type: ignore
            self._stack.append(other)
        else:
            if self._stack:
                self._stack[-1].child_layer = other[0]  # type: ignore
            self._stack.extend(other._stack)
        return self
