import asyncio
from typing import AsyncIterator

import pytest
from asgiref.testing import ApplicationCommunicator
from django.core.asgi import get_asgi_application
from django.http import HttpRequest, StreamingHttpResponse
from django.test import AsyncRequestFactory  # type: ignore
from django.urls import path


async def foobar_stream(sleep_time: float = 1) -> AsyncIterator[bytes]:
    yield b"foo\n"
    await asyncio.sleep(sleep_time)
    yield b"bar\n"


async def foobar_view(request: HttpRequest) -> StreamingHttpResponse:
    """Send a foo to the client, then wait 1s, then send a bar."""
    sleep_time = float(request.GET.get("sleep", 1))
    return StreamingHttpResponse(foobar_stream(sleep_time=sleep_time))


# URL patterns for the test ASGI application
urlpatterns = [
    path("foobar/", foobar_view),
]


@pytest.mark.asyncio
@pytest.mark.urls(__name__)
async def test_foobar_view() -> None:
    """Test that the foobar_view returns b'foo\nbar\n' when called with
    sleep_time=0.001."""
    factory = AsyncRequestFactory()
    scope = factory._base_scope(path="/foobar/", query_string=b"sleep=0.001")
    application = get_asgi_application()
    communicator = ApplicationCommunicator(application, scope)
    await communicator.send_input({"type": "http.request"})

    # fetch headers
    await communicator.receive_output(timeout=1)
    # fetch the 'foo'
    first_response = await communicator.receive_output(timeout=1)
    # fetch the 'bar'
    second_response = await communicator.receive_output(timeout=1)
    # assert first_response["body"], second_response["body"] == (b"foo\n", b"bar\n")
    assert (first_response["body"], second_response["body"]) == (b"foo\n", b"bar\n")

    # fetch the rest of the response so that coroutines are cleaned up
    await communicator.receive_output(timeout=1)
    with pytest.raises(asyncio.TimeoutError):
        await communicator.receive_output(timeout=1)


@pytest.mark.asyncio
@pytest.mark.urls(__name__)
async def test_foobar_view_disconnect() -> None:
    """Test that if we call foobar_view with some large sleep time and then
    disconnect after the first message then the view terminates immediately,
    before sending the 'bar'."""
    factory = AsyncRequestFactory()
    scope = factory._base_scope(path="/foobar/", query_string=b"sleep=0.1")
    application = get_asgi_application()
    communicator = ApplicationCommunicator(application, scope)
    await communicator.send_input({"type": "http.request"})
    await communicator.receive_output(timeout=1)
    first_response = await communicator.receive_output(timeout=1)
    assert first_response["body"] == b"foo\n"
    # simulate disconnection
    await communicator.send_input({"type": "http.disconnect"})
    # ensure that 'bar\n' isn't sent
    with pytest.raises(asyncio.TimeoutError):
        await communicator.receive_output(timeout=0.2)  # type: ignore
