import asyncio
from contextlib import asynccontextmanager


class ServerResponsePartSerializationHelper:
    """
    Part processing serialization for non-cacheable data streams.
    """

    __slots__ = ("_name", "_cond", "_part", "_timeout")

    def __init__(self, name: str, timeout: float = 5.0):
        self._name = name
        self._cond = asyncio.Condition()
        self._part = 0
        self._timeout = timeout

    @asynccontextmanager
    async def serialize_processing(self, part: int):
        async with self._cond:
            try:
                while True:  # Serialize parallel query buffer parts.
                    if part == self._part:
                        break
                    elif part > self._part:
                        async with asyncio.timeout(self._timeout):
                            await self._cond.wait()
                    else:
                        raise Exception(f"Duplicate request for {self._name} part {part}")
                self._part += 1
                yield
            finally:
                self._cond.notify_all()
