import abc
from typing import Callable, Optional, Type, Union

import h11
from h11._readers import ChunkedReader, ContentLengthReader, Http10Reader
from h11._receivebuffer import ReceiveBuffer

from mitmproxy import http, version
from mitmproxy.connection import Connection, ConnectionState
from mitmproxy.net.http import http1, status_codes
from mitmproxy.proxy import commands, events, layer
from mitmproxy.proxy.layers.http._base import ReceiveHttp, StreamId
from mitmproxy.proxy.utils import expect
from mitmproxy.utils import human
from ._base import HttpConnection, format_error
from ._events import HttpEvent, RequestData, RequestEndOfMessage, RequestHeaders, RequestProtocolError, ResponseData, \
    ResponseEndOfMessage, ResponseHeaders, ResponseProtocolError
from ...context import Context

TBodyReader = Union[ChunkedReader, Http10Reader, ContentLengthReader]


class Http1Connection(HttpConnection, metaclass=abc.ABCMeta):
    stream_id: Optional[StreamId] = None
    request: Optional[http.Request] = None
    response: Optional[http.Response] = None
    request_done: bool = False
    response_done: bool = False
    # this is a bit of a hack to make both mypy and PyCharm happy.
    state: Union[Callable[[events.Event], layer.CommandGenerator[None]], Callable]
    body_reader: TBodyReader
    buf: ReceiveBuffer

    ReceiveProtocolError: Type[Union[RequestProtocolError, ResponseProtocolError]]
    ReceiveData: Type[Union[RequestData, ResponseData]]
    ReceiveEndOfMessage: Type[Union[RequestEndOfMessage, ResponseEndOfMessage]]

    def __init__(self, context: Context, conn: Connection):
        super().__init__(context, conn)
        self.buf = ReceiveBuffer()

    @abc.abstractmethod
    def send(self, event: HttpEvent) -> layer.CommandGenerator[None]:
        yield from ()  # pragma: no cover

    @abc.abstractmethod
    def read_headers(self, event: events.ConnectionEvent) -> layer.CommandGenerator[None]:
        yield from ()  # pragma: no cover

    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:
        if isinstance(event, HttpEvent):
            yield from self.send(event)
        else:
            if isinstance(event, events.DataReceived) and self.state != self.passthrough:
                self.buf += event.data
            yield from self.state(event)

    @expect(events.Start)
    def start(self, _) -> layer.CommandGenerator[None]:
        self.state = self.read_headers
        yield from ()

    state = start

    def read_body(self, event: events.Event) -> layer.CommandGenerator[None]:
        assert self.stream_id
        while True:
            try:
                if isinstance(event, events.DataReceived):
                    h11_event = self.body_reader(self.buf)
                elif isinstance(event, events.ConnectionClosed):
                    h11_event = self.body_reader.read_eof()
                else:
                    raise AssertionError(f"Unexpected event: {event}")
            except h11.ProtocolError as e:
                yield commands.CloseConnection(self.conn)
                yield ReceiveHttp(self.ReceiveProtocolError(self.stream_id, f"HTTP/1 protocol error: {e}"))
                return

            if h11_event is None:
                return
            elif isinstance(h11_event, h11.Data):
                data: bytes = bytes(h11_event.data)
                if data:
                    yield ReceiveHttp(self.ReceiveData(self.stream_id, data))
            elif isinstance(h11_event, h11.EndOfMessage):
                assert self.request
                if h11_event.headers:
                    raise NotImplementedError(f"HTTP trailers are not implemented yet.")
                if self.request.data.method.upper() != b"CONNECT":
                    yield ReceiveHttp(self.ReceiveEndOfMessage(self.stream_id))
                is_request = isinstance(self, Http1Server)
                yield from self.mark_done(
                    request=is_request,
                    response=not is_request
                )
                return

    def wait(self, event: events.Event) -> layer.CommandGenerator[None]:
        """
        We wait for the current flow to be finished before parsing the next message,
        as we may want to upgrade to WebSocket or plain TCP before that.
        """
        assert self.stream_id
        if isinstance(event, events.DataReceived):
            return
        elif isinstance(event, events.ConnectionClosed):
            # for practical purposes, we assume that a peer which sent at least a FIN
            # is not interested in any more data from us, see
            # see https://github.com/httpwg/http-core/issues/22
            if event.connection.state is not ConnectionState.CLOSED:
                yield commands.CloseConnection(event.connection)
            yield ReceiveHttp(self.ReceiveProtocolError(self.stream_id, f"Client disconnected.",
                                                        code=status_codes.CLIENT_CLOSED_REQUEST))
        else:  # pragma: no cover
            raise AssertionError(f"Unexpected event: {event}")

    def done(self, event: events.ConnectionEvent) -> layer.CommandGenerator[None]:
        yield from ()  # pragma: no cover

    def make_pipe(self) -> layer.CommandGenerator[None]:
        self.state = self.passthrough
        if self.buf:
            already_received = self.buf.maybe_extract_at_most(len(self.buf))
            yield from self.state(events.DataReceived(self.conn, already_received))

    def passthrough(self, event: events.Event) -> layer.CommandGenerator[None]:
        assert self.stream_id
        if isinstance(event, events.DataReceived):
            yield ReceiveHttp(self.ReceiveData(self.stream_id, event.data))
        elif isinstance(event, events.ConnectionClosed):
            if isinstance(self, Http1Server):
                yield ReceiveHttp(RequestEndOfMessage(self.stream_id))
            else:
                yield ReceiveHttp(ResponseEndOfMessage(self.stream_id))

    def mark_done(self, *, request: bool = False, response: bool = False) -> layer.CommandGenerator[None]:
        if request:
            self.request_done = True
        if response:
            self.response_done = True
        if self.request_done and self.response_done:
            assert self.request
            assert self.response
            if should_make_pipe(self.request, self.response):
                yield from self.make_pipe()
                return
            try:
                read_until_eof_semantics = http1.expected_http_body_size(self.request, self.response) == -1
            except ValueError:
                # this may raise only now (and not earlier) because an addon set invalid headers,
                # in which case it's not really clear what we are supposed to do.
                read_until_eof_semantics = False
            connection_done = (
                read_until_eof_semantics
                or http1.connection_close(self.request.http_version, self.request.headers)
                or http1.connection_close(self.response.http_version, self.response.headers)
                # If we proxy HTTP/2 to HTTP/1, we only use upstream connections for one request.
                # This simplifies our connection management quite a bit as we can rely on
                # the proxyserver's max-connection-per-server throttling.
                or (self.request.is_http2 and isinstance(self, Http1Client))
            )
            if connection_done:
                yield commands.CloseConnection(self.conn)
                self.state = self.done
                return
            self.request_done = self.response_done = False
            self.request = self.response = None
            if isinstance(self, Http1Server):
                self.stream_id += 2
            else:
                self.stream_id = None
            self.state = self.read_headers
            if self.buf:
                yield from self.state(events.DataReceived(self.conn, b""))


class Http1Server(Http1Connection):
    """A simple HTTP/1 server with no pipelining support."""

    ReceiveProtocolError = RequestProtocolError
    ReceiveData = RequestData
    ReceiveEndOfMessage = RequestEndOfMessage
    stream_id: int

    def __init__(self, context: Context):
        super().__init__(context, context.client)
        self.stream_id = 1

    def send(self, event: HttpEvent) -> layer.CommandGenerator[None]:
        assert event.stream_id == self.stream_id
        if isinstance(event, ResponseHeaders):
            self.response = response = event.response

            if response.is_http2:
                response = response.copy()
                # Convert to an HTTP/1 response.
                response.http_version = "HTTP/1.1"
                # not everyone supports empty reason phrases, so we better make up one.
                response.reason = status_codes.RESPONSES.get(response.status_code, "")
                # Shall we set a Content-Length header here if there is none?
                # For now, let's try to modify as little as possible.

            raw = http1.assemble_response_head(response)
            yield commands.SendData(self.conn, raw)
        elif isinstance(event, ResponseData):
            assert self.response
            if "chunked" in self.response.headers.get("transfer-encoding", "").lower():
                raw = b"%x\r\n%s\r\n" % (len(event.data), event.data)
            else:
                raw = event.data
            if raw:
                yield commands.SendData(self.conn, raw)
        elif isinstance(event, ResponseEndOfMessage):
            assert self.response
            if "chunked" in self.response.headers.get("transfer-encoding", "").lower():
                yield commands.SendData(self.conn, b"0\r\n\r\n")
            yield from self.mark_done(response=True)
        elif isinstance(event, ResponseProtocolError):
            if not self.response and event.code != status_codes.NO_RESPONSE:
                yield commands.SendData(self.conn, make_error_response(event.code, event.message))
            if self.conn.state & ConnectionState.CAN_WRITE:
                yield commands.CloseConnection(self.conn)
        else:
            raise AssertionError(f"Unexpected event: {event}")

    def read_headers(self, event: events.ConnectionEvent) -> layer.CommandGenerator[None]:
        if isinstance(event, events.DataReceived):
            request_head = self.buf.maybe_extract_lines()
            if request_head:
                request_head = [bytes(x) for x in request_head]  # TODO: Make url.parse compatible with bytearrays
                try:
                    self.request = http1.read_request_head(request_head)
                    expected_body_size = http1.expected_http_body_size(self.request)
                except ValueError as e:
                    yield commands.Log(f"{human.format_address(self.conn.peername)}: {e}")
                    yield commands.SendData(self.conn, make_error_response(400, str(e)))
                    yield commands.CloseConnection(self.conn)
                    self.state = self.done
                    return
                yield ReceiveHttp(RequestHeaders(self.stream_id, self.request, expected_body_size == 0))
                self.body_reader = make_body_reader(expected_body_size)
                self.state = self.read_body
                yield from self.state(event)
            else:
                pass  # FIXME: protect against header size DoS
        elif isinstance(event, events.ConnectionClosed):
            buf = bytes(self.buf)
            if buf.strip():
                yield commands.Log(f"Client closed connection before completing request headers: {buf!r}")
            yield commands.CloseConnection(self.conn)
        else:
            raise AssertionError(f"Unexpected event: {event}")

    def mark_done(self, *, request: bool = False, response: bool = False) -> layer.CommandGenerator[None]:
        yield from super().mark_done(request=request, response=response)
        if self.request_done and not self.response_done:
            self.state = self.wait


class Http1Client(Http1Connection):
    """A simple HTTP/1 client with no pipelining support."""

    ReceiveProtocolError = ResponseProtocolError
    ReceiveData = ResponseData
    ReceiveEndOfMessage = ResponseEndOfMessage

    def __init__(self, context: Context):
        super().__init__(context, context.server)

    def send(self, event: HttpEvent) -> layer.CommandGenerator[None]:
        if isinstance(event, RequestProtocolError):
            yield commands.CloseConnection(self.conn)
            return

        if not self.stream_id:
            assert isinstance(event, RequestHeaders)
            self.stream_id = event.stream_id
            self.request = event.request
        assert self.stream_id == event.stream_id

        if isinstance(event, RequestHeaders):
            request = event.request
            if request.is_http2:
                # Convert to an HTTP/1 request.
                request = request.copy()  # (we could probably be a bit more efficient here.)
                request.http_version = "HTTP/1.1"
                if "Host" not in request.headers and request.authority:
                    request.headers.insert(0, "Host", request.authority)
                request.authority = ""
            raw = http1.assemble_request_head(request)
            yield commands.SendData(self.conn, raw)
        elif isinstance(event, RequestData):
            assert self.request
            if "chunked" in self.request.headers.get("transfer-encoding", "").lower():
                raw = b"%x\r\n%s\r\n" % (len(event.data), event.data)
            else:
                raw = event.data
            if raw:
                yield commands.SendData(self.conn, raw)
        elif isinstance(event, RequestEndOfMessage):
            assert self.request
            if "chunked" in self.request.headers.get("transfer-encoding", "").lower():
                yield commands.SendData(self.conn, b"0\r\n\r\n")
            elif http1.expected_http_body_size(self.request, self.response) == -1:
                yield commands.CloseConnection(self.conn, half_close=True)
            yield from self.mark_done(request=True)
        else:
            raise AssertionError(f"Unexpected event: {event}")

    def read_headers(self, event: events.ConnectionEvent) -> layer.CommandGenerator[None]:
        if isinstance(event, events.DataReceived):
            if not self.request:
                # we just received some data for an unknown request.
                yield commands.Log(f"Unexpected data from server: {bytes(self.buf)!r}")
                yield commands.CloseConnection(self.conn)
                return
            assert self.stream_id

            response_head = self.buf.maybe_extract_lines()
            if response_head:
                response_head = [bytes(x) for x in response_head]  # TODO: Make url.parse compatible with bytearrays
                try:
                    self.response = http1.read_response_head(response_head)
                    expected_size = http1.expected_http_body_size(self.request, self.response)
                except ValueError as e:
                    yield commands.CloseConnection(self.conn)
                    yield ReceiveHttp(ResponseProtocolError(self.stream_id, f"Cannot parse HTTP response: {e}"))
                    return
                yield ReceiveHttp(ResponseHeaders(self.stream_id, self.response, expected_size == 0))
                self.body_reader = make_body_reader(expected_size)

                self.state = self.read_body
                yield from self.state(event)
            else:
                pass  # FIXME: protect against header size DoS
        elif isinstance(event, events.ConnectionClosed):
            if self.conn.state & ConnectionState.CAN_WRITE:
                yield commands.CloseConnection(self.conn)
            if self.stream_id:
                if self.buf:
                    yield ReceiveHttp(ResponseProtocolError(self.stream_id,
                                                            f"unexpected server response: {bytes(self.buf)!r}"))
                else:
                    # The server has closed the connection to prevent us from continuing.
                    # We need to signal that to the stream.
                    # https://tools.ietf.org/html/rfc7231#section-6.5.11
                    yield ReceiveHttp(ResponseProtocolError(self.stream_id, "server closed connection"))
            else:
                return
        else:
            raise AssertionError(f"Unexpected event: {event}")


def should_make_pipe(request: http.Request, response: http.Response) -> bool:
    if response.status_code == 101:
        return True
    elif response.status_code == 200 and request.method.upper() == "CONNECT":
        return True
    else:
        return False


def make_body_reader(expected_size: Optional[int]) -> TBodyReader:
    if expected_size is None:
        return ChunkedReader()
    elif expected_size == -1:
        return Http10Reader()
    else:
        return ContentLengthReader(expected_size)


def make_error_response(
    status_code: int,
    message: str = "",
) -> bytes:
    resp = http.Response.make(
        status_code,
        format_error(status_code, message),
        http.Headers(
            Server=version.MITMPROXY,
            Connection="close",
            Content_Type="text/html",
        )
    )
    return http1.assemble_response(resp)


__all__ = [
    "Http1Client",
    "Http1Server",
]
