import asyncio
import pickle
import sys
from typing import Generator, Tuple
from unittest.mock import MagicMock, patch

import pytest
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware

from ray._common.utils import get_or_create_event_loop
from ray.serve import HTTPOptions
from ray.serve._private.http_util import (
    ASGIReceiveProxy,
    MessageQueue,
    configure_http_middlewares,
    configure_http_options_with_defaults,
)


@pytest.mark.asyncio
async def test_message_queue_nowait():
    queue = MessageQueue()

    # Check that wait_for_message hangs until a message is sent.
    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(queue.wait_for_message(), 0.001)

    assert len(list(queue.get_messages_nowait())) == 0

    await queue({"type": "http.response.start"})
    await queue.wait_for_message()
    assert len(list(queue.get_messages_nowait())) == 1

    # Check that messages are cleared after being consumed.
    assert len(list(queue.get_messages_nowait())) == 0
    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(queue.wait_for_message(), 0.001)

    # Check that consecutive messages are returned in order.
    await queue({"type": "http.response.start", "idx": 0})
    await queue({"type": "http.response.start", "idx": 1})
    await queue.wait_for_message()
    messages = list(queue.get_messages_nowait())
    assert len(messages) == 2
    assert messages[0]["idx"] == 0
    assert messages[1]["idx"] == 1

    assert len(list(queue.get_messages_nowait())) == 0
    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(queue.wait_for_message(), 0.001)

    # Check that a concurrent waiter is notified when a message is available.
    loop = asyncio.get_running_loop()
    waiting_task = loop.create_task(queue.wait_for_message())
    for _ in range(1000):
        assert not waiting_task.done()

    await queue({"type": "http.response.start"})
    await waiting_task
    assert len(list(queue.get_messages_nowait())) == 1

    # Check that once the queue is closed, new messages should be rejected and
    # ongoing and subsequent calls to wait for messages should return immediately.
    waiting_task = loop.create_task(queue.wait_for_message())
    queue.close()
    await waiting_task  # Ongoing call should return.

    for _ in range(100):
        with pytest.raises(RuntimeError):
            await queue({"hello": "world"})
        await queue.wait_for_message()
        assert queue.get_messages_nowait() == []


@pytest.mark.asyncio
async def test_message_queue_wait():
    queue = MessageQueue()

    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(queue.get_one_message(), 0.001)

    queue.put_nowait("A")
    assert await queue.get_one_message() == "A"

    # Check that messages are cleared after being consumed.
    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(queue.get_one_message(), 0.001)

    # Check that consecutive messages are returned in order.
    queue.put_nowait("B")
    queue.put_nowait("C")
    assert await queue.get_one_message() == "B"
    assert await queue.get_one_message() == "C"

    # Check that messages are cleared after being consumed.
    with pytest.raises(asyncio.TimeoutError):
        await asyncio.wait_for(queue.get_one_message(), 0.001)

    # Check that a concurrent waiter is notified when a message is available.
    loop = asyncio.get_running_loop()
    fetch_task = loop.create_task(queue.get_one_message())
    for _ in range(1000):
        assert not fetch_task.done()
    queue.put_nowait("D")
    assert await fetch_task == "D"


@pytest.mark.asyncio
async def test_message_queue_wait_closed():
    queue = MessageQueue()

    queue.put_nowait("A")
    assert await queue.get_one_message() == "A"

    # Check that once the queue is closed, ongoing and subsequent calls
    # to get_one_message should raise an exception
    loop = asyncio.get_running_loop()
    fetch_task = loop.create_task(queue.get_one_message())
    queue.close()
    with pytest.raises(StopAsyncIteration):
        await fetch_task

    for _ in range(10):
        with pytest.raises(StopAsyncIteration):
            await queue.get_one_message()


@pytest.mark.asyncio
async def test_message_queue_wait_error():
    queue = MessageQueue()

    queue.put_nowait("A")
    assert await queue.get_one_message() == "A"

    # Check setting an error
    loop = asyncio.get_running_loop()
    fetch_task = loop.create_task(queue.get_one_message())
    queue.set_error(TypeError("uh oh! something went wrong."))
    with pytest.raises(TypeError, match="uh oh! something went wrong"):
        await fetch_task

    for _ in range(10):
        with pytest.raises(TypeError, match="uh oh! something went wrong"):
            await queue.get_one_message()


@pytest.fixture
@pytest.mark.asyncio
def setup_receive_proxy(
    request,
) -> Generator[Tuple[ASGIReceiveProxy, MessageQueue], None, None]:
    # Param can be 'http' (default) or 'websocket' (ASGI scope type).
    type = getattr(request, "param", "http")

    queue = MessageQueue()

    async def receive_asgi_messages(request_id: str) -> bytes:
        await queue.wait_for_message()
        messages = queue.get_messages_nowait()
        for message in messages:
            if isinstance(message, Exception):
                raise message

        return pickle.dumps(messages)

    loop = get_or_create_event_loop()
    asgi_receive_proxy = ASGIReceiveProxy({"type": type}, "", receive_asgi_messages)
    receiver_task = loop.create_task(asgi_receive_proxy.fetch_until_disconnect())
    try:
        yield asgi_receive_proxy, queue
    except Exception:
        receiver_task.cancel()


@pytest.mark.asyncio
class TestASGIReceiveProxy:
    async def test_basic(
        self, setup_receive_proxy: Tuple[ASGIReceiveProxy, MessageQueue]
    ):
        asgi_receive_proxy, queue = setup_receive_proxy

        queue.put_nowait({"type": "foo"})
        queue.put_nowait({"type": "bar"})
        assert await asgi_receive_proxy() == {"type": "foo"}
        assert await asgi_receive_proxy() == {"type": "bar"}

        assert asgi_receive_proxy._queue.empty()

        # Once disconnect is received, it should be returned repeatedly.
        queue.put_nowait({"type": "http.disconnect"})
        for _ in range(100):
            assert await asgi_receive_proxy() == {"type": "http.disconnect"}

        # Subsequent messages should be ignored.
        queue.put_nowait({"type": "baz"})
        assert await asgi_receive_proxy() == {"type": "http.disconnect"}

    async def test_raises_exception(
        self, setup_receive_proxy: Tuple[ASGIReceiveProxy, MessageQueue]
    ):
        asgi_receive_proxy, queue = setup_receive_proxy

        queue.put_nowait({"type": "foo"})
        queue.put_nowait({"type": "bar"})
        assert await asgi_receive_proxy() == {"type": "foo"}
        assert await asgi_receive_proxy() == {"type": "bar"}

        queue.put_nowait(RuntimeError("oopsies"))
        with pytest.raises(RuntimeError, match="oopsies"):
            await asgi_receive_proxy()

    @pytest.mark.parametrize(
        "setup_receive_proxy",
        ["http", "websocket"],
        indirect=True,
    )
    async def test_return_disconnect_on_key_error(
        self, setup_receive_proxy: Tuple[ASGIReceiveProxy, MessageQueue]
    ):
        """If the proxy is no longer handling a given request, it raises a KeyError.

        In these cases, the ASGI receive proxy should return a disconnect message.

        See https://github.com/ray-project/ray/pull/44647 for details.
        """
        asgi_receive_proxy, queue = setup_receive_proxy

        queue.put_nowait({"type": "foo"})
        queue.put_nowait({"type": "bar"})
        assert await asgi_receive_proxy() == {"type": "foo"}
        assert await asgi_receive_proxy() == {"type": "bar"}

        queue.put_nowait(KeyError("not found"))
        for _ in range(100):
            if asgi_receive_proxy._type == "http":
                assert await asgi_receive_proxy() == {"type": "http.disconnect"}
            else:
                assert await asgi_receive_proxy() == {
                    "type": "websocket.disconnect",
                    "code": 1005,
                }

    async def test_receive_asgi_messages_raises(self):
        async def receive_asgi_messages(request_id: str) -> bytes:
            raise RuntimeError("maybe actor crashed")

        loop = get_or_create_event_loop()
        asgi_receive_proxy = ASGIReceiveProxy(
            {"type": "http"}, "", receive_asgi_messages
        )
        receiver_task = loop.create_task(asgi_receive_proxy.fetch_until_disconnect())

        try:
            with pytest.raises(RuntimeError, match="maybe actor crashed"):
                await asgi_receive_proxy()
        finally:
            receiver_task.cancel()


class MockMiddleware:
    """Mock middleware class for testing."""

    def __init__(self, name):
        self.name = name

    def __eq__(self, other):
        return isinstance(other, MockMiddleware) and self.name == other.name

    def __repr__(self):
        return f"MockMiddleware({self.name})"


@pytest.fixture
def base_http_options():
    """Provides basic HTTPOptions for testing."""
    return HTTPOptions(
        host="0.0.0.0",
        port=8000,
        request_timeout_s=30.0,
        keep_alive_timeout_s=5.0,
        middlewares=[],
    )


class TestConfigureHttpOptionsWithDefaults:
    """Test suite for configure_http_options_with_defaults function."""

    def test_basic_configuration(self, base_http_options):
        """Test basic configuration preserves settings."""
        result = configure_http_options_with_defaults(base_http_options)

        # Request timeout should be preserved
        assert result.request_timeout_s == 30.0
        # Keep alive timeout should be preserved (no env override)
        assert result.keep_alive_timeout_s == 5.0
        # Should initialize middlewares list
        assert result.middlewares == []
        # Original should not be modified
        assert base_http_options.request_timeout_s == 30.0

    @patch("ray.serve._private.http_util.call_function_from_import_path")
    @patch(
        "ray.serve._private.http_util.RAY_SERVE_HTTP_PROXY_CALLBACK_IMPORT_PATH",
        "my.module.callback",
    )
    def test_callback_middleware_injection(self, mock_call_function, base_http_options):
        """Test that the callback middleware is injected correctly."""

        # Arrange: Create a valid middleware by wrapping it with Starlette's Middleware class
        class CustomMiddleware(BaseHTTPMiddleware):
            async def dispatch(self, request, call_next):
                response = await call_next(request)  # Simply pass the request through
                return response

        # Mock the app argument
        mock_app = MagicMock()

        wrapped_middleware = Middleware(CustomMiddleware, app=mock_app)
        mock_call_function.return_value = [
            wrapped_middleware
        ]  # Return list of wrapped middleware

        # Act
        result = configure_http_middlewares(base_http_options)

        # Assert
        mock_call_function.assert_called_once_with(
            "my.module.callback"
        )  # Verify callback execution
        assert len(result.middlewares) == 1  # Ensure one middleware was injected
        assert isinstance(result.middlewares[0], Middleware)

    def test_callback_middleware_disabled(self, base_http_options):
        """Test that callback middleware is not loaded when disabled."""
        with patch(
            "ray.serve._private.http_util.RAY_SERVE_HTTP_PROXY_CALLBACK_IMPORT_PATH",
            "",
        ):
            result = configure_http_options_with_defaults(base_http_options)

            # Assert that no callback middleware is added
            assert result.middlewares == []

    def test_deep_copy_behavior(self, base_http_options):
        """Test that an original HTTPOptions object is not modified."""
        original_timeout = base_http_options.request_timeout_s

        result = configure_http_options_with_defaults(base_http_options)

        # Original should remain unchanged
        assert base_http_options.request_timeout_s == original_timeout
        # Result should be a different object
        assert result is not base_http_options


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", "-s", __file__]))