from typing import TYPE_CHECKING, List, Tuple, Union
from unittest.mock import MagicMock
from warnings import catch_warnings

import pytest

from litestar import MediaType, asgi, get
from litestar.datastructures.headers import MutableScopeHeaders
from litestar.exceptions import LitestarWarning, ValidationException
from litestar.middleware import AbstractMiddleware, ASGIMiddleware, DefineMiddleware
from litestar.params import FromPath
from litestar.response.base import ASGIResponse
from litestar.status_codes import HTTP_400_BAD_REQUEST
from litestar.testing import create_test_client

if TYPE_CHECKING:
    from litestar.types import ASGIApp, Message, Receive, Scope, Send


def test_custom_middleware() -> None:
    class SubclassMiddleware(AbstractMiddleware):
        async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

            await self.app(scope, receive, _send)

    @get("/")
    def handler() -> dict:
        return {"hello": "world"}

    with create_test_client(handler, middleware=[DefineMiddleware(SubclassMiddleware)]) as client:
        response = client.get("/")
        assert response.headers["test"] == "123"


def test_raises_exception() -> None:
    class SubclassMiddleware(AbstractMiddleware):
        async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
            raise ValidationException(detail="nope")

    @get("/")
    def handler() -> dict:
        return {"hello": "world"}

    with create_test_client(handler, middleware=[DefineMiddleware(SubclassMiddleware)]) as client:
        response = client.get("/")
        assert response.status_code == HTTP_400_BAD_REQUEST


def test_exclude_by_pattern() -> None:
    class SubclassMiddleware(AbstractMiddleware):
        exclude = r"^/123"

        async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

            await self.app(scope, receive, _send)

    @get("/123")
    def first_handler() -> dict:
        return {"hello": "world"}

    @get("/456")
    def second_handler() -> dict:
        return {"hello": "world"}

    @asgi("/mount", is_mount=True)
    async def handler(scope: "Scope", receive: "Receive", send: "Send") -> None:
        response = ASGIResponse(body=b"ok", media_type=MediaType.TEXT)
        await response(scope, receive, send)

    with create_test_client(
        [first_handler, second_handler, handler], middleware=[DefineMiddleware(SubclassMiddleware)]
    ) as client:
        response = client.get("/123")
        assert "test" not in response.headers

        response = client.get("/456")
        assert "test" in response.headers

        response = client.get("/mount/123")
        assert "test" in response.headers


def test_exclude_by_pattern_list() -> None:
    class SubclassMiddleware(AbstractMiddleware):
        exclude = ["123", "456"]

        async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

            await self.app(scope, receive, _send)

    @get("/123")
    def first_handler() -> dict:
        return {"hello": "world"}

    @get("/456")
    def second_handler() -> dict:
        return {"hello": "world"}

    @get("/789")
    def third_handler() -> dict:
        return {"hello": "world"}

    with create_test_client(
        [first_handler, second_handler, third_handler], middleware=[DefineMiddleware(SubclassMiddleware)]
    ) as client:
        response = client.get("/123")
        assert "test" not in response.headers
        response = client.get("/456")
        assert "test" not in response.headers
        response = client.get("/789")
        assert "test" in response.headers


@pytest.mark.parametrize("excludes", ["/", ["/", "/foo"], "/*", "/.*"])
def test_exclude_by_pattern_warns_if_exclude_all(excludes: Union[str, List[str]]) -> None:
    class SubclassMiddleware(AbstractMiddleware):
        exclude = excludes

        async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
            await self.app(scope, receive, send)

    with pytest.warns(LitestarWarning, match="Middleware 'SubclassMiddleware' exclude pattern"):
        create_test_client(middleware=[SubclassMiddleware])


def test_exclude_doesnt_warn_on_non_greedy_pattern() -> None:
    class SubclassMiddleware(AbstractMiddleware):
        exclude = "^/$"

        async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
            await self.app(scope, receive, send)

    with catch_warnings(record=True) as warnings:
        create_test_client(middleware=[SubclassMiddleware])
        assert len(warnings) == 0


def test_exclude_by_opt_key() -> None:
    class SubclassMiddleware(AbstractMiddleware):
        exclude_opt_key = "exclude_route"

        async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

                await self.app(scope, receive, _send)

    @get("/", exclude_route=True)
    def handler() -> dict:
        return {"hello": "world"}

    with create_test_client(handler, middleware=[DefineMiddleware(SubclassMiddleware)]) as client:
        response = client.get("/")
        assert "test" not in response.headers


def test_abstract_middleware_deprecation_warning() -> None:
    with pytest.warns(DeprecationWarning, match="AbstractMiddleware"):

        class MyMiddleware(AbstractMiddleware):
            pass


def test_asgi_middleware() -> None:
    class SubclassMiddleware(ASGIMiddleware):
        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

            await next_app(scope, receive, _send)

    @get("/")
    def handler() -> dict:
        return {"hello": "world"}

    with create_test_client(handler, middleware=[SubclassMiddleware()]) as client:
        response = client.get("/")
        assert response.headers["test"] == "123"


def test_asgi_middleware_raises_exception() -> None:
    class SubclassMiddleware(ASGIMiddleware):
        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            raise ValidationException(detail="nope")

    @get("/")
    def handler() -> dict:
        return {"hello": "world"}

    with create_test_client(handler, middleware=[SubclassMiddleware()]) as client:
        response = client.get("/")
        assert response.status_code == HTTP_400_BAD_REQUEST


def test_asgi_middleware_exclude_by_pattern() -> None:
    class SubclassMiddleware(ASGIMiddleware):
        def __init__(self) -> None:
            self.exclude_path_pattern = r"^/123"

        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

            await next_app(scope, receive, _send)

    @get("/123")
    def first_handler() -> dict:
        return {"hello": "world"}

    @get("/456")
    def second_handler() -> dict:
        return {"hello": "world"}

    @asgi("/mount", is_mount=True)
    async def handler(scope: "Scope", receive: "Receive", send: "Send") -> None:
        response = ASGIResponse(body=b"ok", media_type=MediaType.TEXT)
        await response(scope, receive, send)

    with create_test_client([first_handler, second_handler, handler], middleware=[SubclassMiddleware()]) as client:
        response = client.get("/123")
        assert "test" not in response.headers

        response = client.get("/456")
        assert "test" in response.headers

        response = client.get("/mount/123")
        assert "test" in response.headers


def test_asgi_middleware_exclude_by_pattern_tuple() -> None:
    class SubclassMiddleware(ASGIMiddleware):
        exclude_path_pattern = ("123", "456")

        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

            await next_app(scope, receive, _send)

    @get("/123")
    def first_handler() -> dict:
        return {"hello": "world"}

    @get("/456")
    def second_handler() -> dict:
        return {"hello": "world"}

    @get("/789")
    def third_handler() -> dict:
        return {"hello": "world"}

    with create_test_client(
        [first_handler, second_handler, third_handler], middleware=[SubclassMiddleware()]
    ) as client:
        response = client.get("/123")
        assert "test" not in response.headers
        response = client.get("/456")
        assert "test" not in response.headers
        response = client.get("/789")
        assert "test" in response.headers


def test_asgi_middleware_should_exclude_scope() -> None:
    mock = MagicMock()

    class SubclassMiddleware(ASGIMiddleware):
        @staticmethod
        def should_bypass_for_scope(scope: "Scope") -> bool:
            return scope["path"].endswith(".jpg")

        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            mock(scope["path"])
            await next_app(scope, receive, send)

    @get("/{file_name:str}")
    def handler(file_name: FromPath[str]) -> str:
        return file_name

    with create_test_client([handler], middleware=[SubclassMiddleware()]) as client:
        assert client.get("/test.txt").status_code == 200
        assert client.get("/test.jpg").status_code == 200

        mock.assert_called_once_with("/test.txt")


def test_asgi_middleware_path_exclude_warns_future_use() -> None:
    mock = MagicMock()

    class SubclassMiddleware(ASGIMiddleware):
        def __init__(self, pattern: str) -> None:
            self.exclude_path_pattern = pattern

        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            mock(scope["path"])
            await next_app(scope, receive, send)

    @get("/{file_name:str}")
    def handler(file_name: FromPath[str]) -> str:
        return file_name

    # this configuration would NOT be excluded in the future
    with create_test_client([handler], middleware=[SubclassMiddleware(".jpg")]) as client:
        with pytest.warns(DeprecationWarning, match=".*exclude_path_pattern.* did match the request path"):
            assert client.get("/test.jpg").status_code == 200

        assert client.get("/test.txt").status_code == 200
        mock.assert_called_once_with("/test.txt")

    mock.reset_mock()

    # this configuration WOULD be excluded in the future
    with create_test_client([handler], middleware=[SubclassMiddleware("str")]) as client:
        with pytest.warns(DeprecationWarning, match=".*exclude_path_pattern.* did not match the request path"):
            assert client.get("/test.jpg").status_code == 200
            mock.assert_called_once_with("/test.jpg")


@pytest.mark.parametrize("excludes", ["/", ("/", "/foo"), "/*", "/.*"])
def test_asgi_middleware_exclude_by_pattern_warns_if_exclude_all(excludes: Union[str, Tuple[str, ...]]) -> None:
    class SubclassMiddleware(ASGIMiddleware):
        exclude_path_pattern = excludes

        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            await next_app(scope, receive, send)

    with pytest.warns(LitestarWarning, match="Middleware 'SubclassMiddleware' exclude pattern"):
        create_test_client(middleware=[SubclassMiddleware()])


def test_asgi_middleware_exclude_doesnt_warn_on_non_greedy_pattern() -> None:
    class SubclassMiddleware(ASGIMiddleware):
        exclude_path_pattern = "^/$"

        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            await next_app(scope, receive, send)

    with catch_warnings(record=True) as warnings:
        create_test_client(middleware=[SubclassMiddleware()])
        assert len(warnings) == 0


def test_asgi_middleware_exclude_by_opt_key() -> None:
    class SubclassMiddleware(ASGIMiddleware):
        exclude_opt_key = "exclude_route"

        async def handle(self, scope: "Scope", receive: "Receive", send: "Send", next_app: "ASGIApp") -> None:
            async def _send(message: "Message") -> None:
                if message["type"] == "http.response.start":
                    headers = MutableScopeHeaders(message)
                    headers.add("test", str(123))
                await send(message)

                await next_app(scope, receive, send)

    @get("/", exclude_route=True)
    def handler() -> dict:
        return {"hello": "world"}

    with create_test_client(handler, middleware=[SubclassMiddleware()]) as client:
        response = client.get("/")
        assert "test" not in response.headers
