"""Bound wire and decoded response bytes before JSON parsing."""
import asyncio
import json
import zlib

from backend.errors import UpstreamFailedError, UpstreamResponseTooLargeError


async def read_json(response, limit):
    try:
        length = response.headers.get("content-length")
        if length is not None and int(length) > limit:
            raise UpstreamResponseTooLargeError()
        if response.is_stream_consumed:
            body = response.content
            if len(body) > limit:
                raise UpstreamResponseTooLargeError()
        else:
            encoding = response.headers.get("content-encoding", "identity").lower()
            if encoding not in ("identity", "gzip", "deflate"):
                raise UpstreamFailedError("Unsupported upstream content encoding")
            decoder = zlib.decompressobj(31 if encoding == "gzip" else 15) if encoding != "identity" else None
            body = bytearray()
            wire_bytes = 0
            async for chunk in response.aiter_raw(chunk_size=64 * 1024):
                wire_bytes += len(chunk)
                if wire_bytes > limit:
                    raise UpstreamResponseTooLargeError()
                decoded = decoder.decompress(chunk, limit - len(body) + 1) if decoder else chunk
                if len(body) + len(decoded) > limit or (decoder and decoder.unconsumed_tail):
                    raise UpstreamResponseTooLargeError()
                body.extend(decoded)
            if decoder:
                tail = decoder.flush(limit - len(body) + 1)
                if len(body) + len(tail) > limit:
                    raise UpstreamResponseTooLargeError()
                body.extend(tail)
                if not decoder.eof or decoder.unused_data:
                    raise UpstreamFailedError("Invalid upstream compressed response")
        return await asyncio.to_thread(json.loads, body)
    except (UpstreamFailedError, UpstreamResponseTooLargeError):
        raise
    except Exception as exc:
        raise UpstreamFailedError("Invalid upstream JSON response") from exc
    finally:
        await response.aclose()