"""Потоковая отдача файла с детерминированным освобождением источника (T9.1, Review Focus 4 плана Среза 3b).

Starlette `StreamingResponse.stream_response` не вызывает `aclose()` у `body_iterator`: при обрыве клиента или если тело
не начали итерировать поток S3 освобождался бы только сборщиком мусора. `ClosingStreamingResponse` закрывает итератор
в `finally`, а `ChainedStream` гарантирует, что `aclose` закрывает и исходный поток, даже если обёртку не начинали
читать (у неначатого асинхронного генератора `finally` не выполняется, поэтому обёртка — объект, а не генератор).
"""

from collections.abc import AsyncGenerator

import anyio
from fastapi.responses import StreamingResponse
from starlette.types import Send


class ChainedStream:
    """Уже прочитанная первая порция и остальной поток `rest`; `aclose` всегда закрывает `rest` (там `close` + `release`
    ответа S3), в том числе при завершении, ошибке и обрыве потребителя, начатом и не начатом чтении."""

    def __init__(self, first: bytes, rest: AsyncGenerator[bytes, None]) -> None:
        self._first = first
        self._rest = rest

    def __aiter__(self) -> "ChainedStream":
        return self

    async def __anext__(self) -> bytes:
        if self._first:
            chunk, self._first = self._first, b""
            return chunk
        return await anext(self._rest)

    async def aclose(self) -> None:
        """Закрыть исходный поток и сбросить префетченную порцию: после закрытия `__anext__` её не отдаёт
        (идемпотентно: у завершённого генератора `aclose` — no-op)."""
        self._first = b""
        await self._rest.aclose()


class ClosingStreamingResponse(StreamingResponse):
    """`StreamingResponse`, закрывающий `body_iterator` (`aclose`) после отдачи, ошибки или обрыва клиента.

    Закрытие идёт под `CancelScope(shield=True)`: при обрыве клиента задача отдачи отменяется, а освобождение
    соединения S3 требует `await`."""

    async def stream_response(self, send: Send) -> None:
        try:
            await super().stream_response(send)
        finally:
            aclose = getattr(self.body_iterator, "aclose", None)
            if aclose is not None:
                with anyio.CancelScope(shield=True):
                    await aclose()