"""Regression tests for LIT-2642 — interrupted streams must still flush usage."""

import asyncio
from typing import List
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest


def _make_streaming_response(chunks: List[bytes]):
    mock = MagicMock(spec=httpx.Response)
    mock.status_code = 200
    mock.headers = httpx.Headers({"content-type": "application/vnd.amazon.eventstream"})
    mock.raise_for_status = MagicMock(return_value=None)

    async def _aiter_bytes():
        for chunk in chunks:
            yield chunk

    mock.aiter_bytes = _aiter_bytes
    mock.aclose = AsyncMock()
    return mock


def _make_logging_obj():
    mock = MagicMock()
    mock.async_flush_passthrough_collected_chunks = AsyncMock()
    return mock


class _ImmediateExecutor:
    def submit(self, fn, *args, **kwargs):
        fn(*args, **kwargs)


class _RecordingCollector:
    def __init__(self) -> None:
        self.chunks: List[bytes] = []

    def add(self, chunk: bytes) -> None:
        self.chunks.append(chunk)

    def build_logged_response(self, litellm_logging_obj: MagicMock) -> bytes:
        return b"".join(self.chunks)


class _FailingCollector(_RecordingCollector):
    def add(self, chunk: bytes) -> None:
        raise ValueError("bad frame")


def _provider_config(collector: _RecordingCollector) -> MagicMock:
    provider_config = MagicMock()
    provider_config.create_stream_collector.return_value = collector
    return provider_config


def _spend_payload(flush_mock: MagicMock) -> bytes:
    flush_mock.assert_called_once()
    collector = flush_mock.call_args.kwargs["collector"]
    return collector.build_logged_response(litellm_logging_obj=MagicMock())


@pytest.mark.asyncio
async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion():
    from litellm.passthrough.main import AsyncPassthroughStreamingResponse

    chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
    mock_response = _make_streaming_response(chunks)
    mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})

    async def response_coro():
        return mock_response

    mock_logging_obj = _make_logging_obj()

    received = []
    received_response = AsyncPassthroughStreamingResponse(
        response=response_coro(),
        litellm_logging_obj=mock_logging_obj,
        provider_config=_provider_config(_RecordingCollector()),
    )

    async for chunk in received_response:
        received.append(chunk)

    assert received == chunks

    assert received_response.headers["content-type"] == "application/octet-stream"
    assert received_response.headers["x-request-id"] == "req-123"

    await asyncio.sleep(0)

    assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == b"".join(chunks)


@pytest.mark.asyncio
async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect():
    from litellm.passthrough.main import AsyncPassthroughStreamingResponse

    chunks = [
        b'{"chunk": 1, "outputTokens": 10}',
        b'{"chunk": 2, "outputTokens": 12}',
        b'{"chunk": 3, "outputTokens": 8}',
    ]
    mock_response = _make_streaming_response(chunks)
    mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})

    async def response_coro():
        return mock_response

    mock_logging_obj = _make_logging_obj()

    gen = AsyncPassthroughStreamingResponse(
        response=response_coro(),
        litellm_logging_obj=mock_logging_obj,
        provider_config=_provider_config(_RecordingCollector()),
    )

    received = [await gen.__anext__()]
    await gen.aclose()

    assert received == [chunks[0]]

    await asyncio.sleep(0)

    assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == chunks[0]


@pytest.mark.asyncio
async def test_asyncpassthroughstreamingresponse_does_not_flush_on_4xx():
    from litellm.passthrough.main import AsyncPassthroughStreamingResponse

    err_response = MagicMock(spec=httpx.Response)
    err_response.status_code = 429
    err_response.headers = httpx.Headers({"content-type": "application/octet-stream"})

    def _raise():
        raise httpx.HTTPStatusError(
            "429",
            request=httpx.Request("POST", "https://example.com"),
            response=httpx.Response(429, request=httpx.Request("POST", "https://example.com")),
        )

    err_response.raise_for_status = _raise
    err_response.aclose = AsyncMock()

    async def response_coro():
        return err_response

    mock_logging_obj = _make_logging_obj()

    with pytest.raises(httpx.HTTPStatusError):
        async for _ in AsyncPassthroughStreamingResponse(
            response=response_coro(),
            litellm_logging_obj=mock_logging_obj,
            provider_config=MagicMock(),
        ):
            pass

    mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()


@pytest.mark.asyncio
async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_with_partial_data():
    from litellm.passthrough.main import AsyncPassthroughStreamingResponse

    partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"]

    mock_response = MagicMock(spec=httpx.Response)
    mock_response.status_code = 200
    mock_response.raise_for_status = MagicMock(return_value=None)
    mock_response.aclose = AsyncMock()
    mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})

    async def _aiter_bytes_then_raise():
        for c in partial_chunks:
            yield c
        raise httpx.ReadError("upstream disconnected")

    mock_response.aiter_bytes = _aiter_bytes_then_raise

    async def response_coro():
        return mock_response

    mock_logging_obj = _make_logging_obj()

    received = []

    async def _drain():
        async for chunk in AsyncPassthroughStreamingResponse(
            response=response_coro(),
            litellm_logging_obj=mock_logging_obj,
            provider_config=_provider_config(_RecordingCollector()),
        ):
            received.append(chunk)

    with pytest.raises(httpx.ReadError):
        await _drain()

    assert received == partial_chunks

    await asyncio.sleep(0)

    assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == b"".join(partial_chunks)


def test_passthroughstreamingresponse_flushes_on_normal_completion():
    from litellm.passthrough.main import PassthroughStreamingResponse

    chunks = [b"a", b"b", b"c"]

    mock_response = MagicMock(spec=httpx.Response)
    mock_response.status_code = 200
    mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})

    def _iter_bytes():
        yield from chunks

    mock_response.iter_bytes = _iter_bytes

    mock_logging_obj = MagicMock()
    mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()

    received_responce = PassthroughStreamingResponse(
        response=mock_response,
        litellm_logging_obj=mock_logging_obj,
        provider_config=_provider_config(_RecordingCollector()),
    )

    with patch("litellm.utils.executor", _ImmediateExecutor()):
        received = list(received_responce)

    assert received == chunks

    assert received_responce.headers["content-type"] == "application/octet-stream"
    assert received_responce.headers["x-request-id"] == "req-123"

    assert _spend_payload(mock_logging_obj.flush_passthrough_collected_chunks) == b"".join(chunks)


def test_passthroughstreamingresponse_flushes_on_early_close():
    from litellm.passthrough.main import PassthroughStreamingResponse

    chunks = [b"first", b"second", b"third"]

    mock_response = MagicMock(spec=httpx.Response)
    mock_response.status_code = 200
    mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})

    def _iter_bytes():
        yield from chunks

    mock_response.iter_bytes = _iter_bytes

    mock_logging_obj = MagicMock()
    mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()

    with patch("litellm.utils.executor", _ImmediateExecutor()):
        gen = PassthroughStreamingResponse(
            response=mock_response,
            litellm_logging_obj=mock_logging_obj,
            provider_config=_provider_config(_RecordingCollector()),
        )

        first = next(gen)
        gen.close()

    assert first == chunks[0]
    assert _spend_payload(mock_logging_obj.flush_passthrough_collected_chunks) == chunks[0]


@pytest.mark.asyncio
async def test_asyncpassthroughstreamingresponse_relays_the_stream_when_spend_parsing_fails():
    from litellm.passthrough.main import AsyncPassthroughStreamingResponse

    chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
    mock_response = _make_streaming_response(chunks)

    async def response_coro():
        return mock_response

    mock_logging_obj = _make_logging_obj()

    received = [
        chunk
        async for chunk in AsyncPassthroughStreamingResponse(
            response=response_coro(),
            litellm_logging_obj=mock_logging_obj,
            provider_config=_provider_config(_FailingCollector()),
        )
    ]
    await asyncio.sleep(0)

    assert received == chunks
    mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()


def test_passthroughstreamingresponse_relays_the_stream_when_spend_parsing_fails():
    from litellm.passthrough.main import PassthroughStreamingResponse

    chunks = [b"a", b"b", b"c"]
    mock_response = MagicMock(spec=httpx.Response)
    mock_response.status_code = 200
    mock_response.headers = httpx.Headers({"content-type": "application/octet-stream"})
    mock_response.iter_bytes = lambda: iter(chunks)

    mock_logging_obj = MagicMock()
    mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()

    received = list(
        PassthroughStreamingResponse(
            response=mock_response,
            litellm_logging_obj=mock_logging_obj,
            provider_config=_provider_config(_FailingCollector()),
        )
    )

    assert received == chunks
    mock_logging_obj.flush_passthrough_collected_chunks.assert_not_called()