"""Bounded optional IMAGE returns, independent of ComfyUI and Distributed.

Exact mode is the existing tensor wire format, unchanged. All other modes clip
and truncate RGB/luma to uint8 (as the Collector does), normalize in float32,
then cast to the source dtype. RGBA always carries an exact tensor alpha
sidecar, including PNG/WebP. ``legacy_png`` means uncompressed binary PNG here,
not the old Collector's JSON/base64 transport.

The compressed wire format concatenates length-delimited stills or one Matroska
video, followed by optional alpha bytes. No archive, path, URL or pickle is
accepted. Limits cover wire bytes, descriptor structure, output tensor bytes,
and aggregate padded pixel buffers, not process RSS/native codec overhead.
Video preflight reads EBML track dimensions and block counts before creating a
decoder. We extract packets ourselves, avoiding native container stream probing
(which can decode before caller checks). FFmpeg max_pixels also bounds
bitstream-driven resolution changes, including native buffer alignment.
Software codecs only, one thread, no fallback to another encoder/format.
"""

import io
import struct
import zlib
from fractions import Fraction

import torch

from . import codec
from .codec import MAX_PAYLOAD_BYTES
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options

MAX_FRAMES = 240
MAX_DIMENSION = 8192
MAX_DESCRIPTOR_BYTES = 64 * 1024
MAX_CONTAINER_ELEMENTS = 16384
_VIDEO_CODECS = {
    "ffv1": ("ffv1", None, "bgr0"),
    "h264": ("libx264", 51, "yuv420p"),
    "h265": ("libx265", 51, "yuv420p"),
    "av1": ("libsvtav1", 63, "yuv420p"),
}
_CODEC_IDS = {"ffv1": b"V_FFV1", "h264": b"V_MPEG4/ISO/AVC",
              "h265": b"V_MPEGH/ISO/HEVC", "av1": b"V_AV1"}


class _BoundedIO(io.BytesIO):
    """Enforce the final file size on every write, including muxer flush/seek."""
    def __init__(self, limit):
        super().__init__()
        self.limit = limit

    def write(self, data):
        if self.tell() + len(data) > self.limit:
            raise ValueError("Encoded IMAGE exceeds byte limit")
        return super().write(data)

    def seek(self, offset, whence=0):
        if whence == 0:
            position = offset
        elif whence == 1:
            position = self.tell() + offset
        elif whence == 2:
            position = self.getbuffer().nbytes + offset
        else:
            raise ValueError("Invalid image buffer seek")
        if not 0 <= position <= self.limit:
            raise ValueError("Encoded IMAGE exceeds byte limit")
        return super().seek(offset, whence)


def _quality_to_crf(quality, maximum):
    return int(round(maximum * (1.0 - quality / 100.0) ** 0.6))


def _padded_size(fmt, height, width):
    if fmt in ("h264", "h265", "av1"):
        height, width = height + height % 2, width + width % 2
    # libx265 / SVT-AV1 cannot safely encode arbitrarily small pictures.
    if fmt in ("h265", "av1"):
        height, width = max(64, height), max(64, width)
    return height, width


def _shape_budget(shape, fmt, limit, itemsize=4, alpha_size=0, coded_frames=None):
    size = codec._tensor_size("IMAGE", shape, itemsize, limit)
    count, height, width, _ = shape
    if count > MAX_FRAMES or max(height, width) > MAX_DIMENSION:
        raise ValueError("IMAGE frame count or dimensions exceed limit")
    ph, pw = _padded_size(fmt, height, width)
    # Reserve output and padded RGBA-sized decoded buffers BEFORE allocation.
    bh, bw = ((ph + 63) // 64 * 64, (pw + 63) // 64 * 64) if fmt in _VIDEO_CODECS else (ph, pw)
    # Slice/OBU counts conservatively bound native frames, including AV1 hidden
    # frames and multiple outputs from one packet. Reserve 16 bytes per pixel
    # for hostile high-bit-depth streams, not only our usual 8-bit RGB frames.
    native_count = max(count, coded_frames or count)
    if native_count > MAX_FRAMES * 4:
        raise ValueError("IMAGE coded frame count exceeds limit")
    pixels = codec._tensor_size("IMAGE", [native_count, bh, bw, 4], 4, limit)
    if size + pixels + alpha_size > limit:
        raise ValueError("IMAGE decoded aggregate exceeds byte limit")
    return ph, pw


def _descriptor_budget(value):
    """Bound the JSON-shaped tree without serializing attacker-sized values."""
    remaining = MAX_DESCRIPTOR_BYTES
    def walk(obj, depth):
        nonlocal remaining
        if depth > 5:
            raise ValueError("IMAGE descriptor nesting exceeds limit")
        remaining -= 8
        if remaining < 0:
            raise ValueError("IMAGE descriptor exceeds byte limit")
        if obj is None:
            return
        if type(obj) is str:
            remaining -= len(obj) * 6
        elif type(obj) is int:
            if obj.bit_length() > 63:
                raise ValueError("IMAGE descriptor integer exceeds limit")
            remaining -= 20
        elif type(obj) is list:
            if len(obj) > MAX_FRAMES:
                raise ValueError("IMAGE descriptor list exceeds limit")
            for item in obj:
                walk(item, depth + 1)
        elif type(obj) is dict:
            if len(obj) > 10:
                raise ValueError("IMAGE descriptor fields exceed limit")
            for key, item in obj.items():
                if type(key) is not str:
                    raise ValueError("Invalid IMAGE descriptor key")
                walk(key, depth + 1)
                walk(item, depth + 1)
        else:
            raise ValueError("Invalid IMAGE descriptor value")
        if remaining < 0:
            raise ValueError("IMAGE descriptor exceeds byte limit")
    walk(value, 0)


def _alpha_metadata(alpha, shape, limit):
    if shape[-1] != 4:
        if alpha is not None:
            raise ValueError("Unexpected IMAGE alpha sidecar")
        return 0, torch.float32
    if type(alpha) is not dict or set(alpha) != {"descriptor", "size"}:
        raise ValueError("Missing or invalid IMAGE alpha sidecar")
    d = alpha["descriptor"]
    if type(d) is not dict or set(d) != {"encoding", "dtype", "shape", "layout"}:
        raise ValueError("Invalid IMAGE alpha descriptor")
    if d["encoding"] != "tensor-le-v1" or d["layout"] != "BHWC" or d["shape"] != shape[:3] + [1]:
        raise ValueError("Invalid IMAGE alpha shape or encoding")
    if type(d["dtype"]) is not str or d["dtype"] not in codec._DTYPES:
        raise ValueError("Invalid IMAGE alpha dtype")
    dtype, itemsize = codec._DTYPES[d["dtype"]]
    size = codec._tensor_size("IMAGE", d["shape"], itemsize, limit)
    if type(alpha["size"]) is not int or alpha["size"] != size:
        raise ValueError("Invalid IMAGE alpha size")
    return size, torch.float64 if dtype == torch.float64 else torch.float32


def _import_av():
    try:
        import av
    except ImportError:
        raise ValueError("PyAV is required for IMAGE video transfer") from None
    return av


def _encode_still(frame, fmt, quality, limit):
    from PIL import Image
    sink = _BoundedIO(limit)
    if frame.shape[-1] == 1:
        frame = frame[..., 0]
    image = Image.fromarray(frame)
    if fmt in ("legacy_png", "png"):
        image.save(sink, format="PNG", compress_level=0 if fmt == "legacy_png" else 4)
    elif fmt == "jpeg":
        image.save(sink, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
    else:
        image.save(sink, format="WEBP", quality=quality, method=4)
    return sink.getvalue()


def _encode_video(frames, fmt, quality, height, width, limit):
    import numpy as np
    av = _import_av()
    encoder, maximum, pix_fmt = _VIDEO_CODECS[fmt]
    try:
        av.codec.Codec(encoder, "w")
    except Exception:
        raise ValueError("IMAGE software encoder unavailable: " + encoder) from None
    sink = _BoundedIO(limit)
    options = {}
    if maximum is not None:
        options["crf"] = str(_quality_to_crf(quality, maximum))
    if fmt == "h264":
        options["preset"] = "medium"
    elif fmt == "h265":
        options.update(preset="fast", **{"x265-params": "log-level=error:pools=none:frame-threads=1"})
    elif fmt == "av1":
        options.update(preset="8", **{"svtav1-params": "lp=1"})
    else:
        options.update(level="3", slicecrc="1")
    try:
        with av.open(sink, mode="w", format="matroska") as container:
            stream = container.add_stream(encoder, rate=24)
            stream.width, stream.height = width, height
            stream.pix_fmt = pix_fmt
            stream.codec_context.thread_count = 1
            stream.options = options
            for index, frame in enumerate(frames):
                if frame.shape[-1] == 1:
                    frame = np.repeat(frame, 3, axis=-1)
                h, w = frame.shape[:2]
                frame = np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
                vf = av.VideoFrame.from_ndarray(frame, format="rgb24")
                vf.pts = index
                vf.time_base = Fraction(1, 24)
                for packet in stream.encode(vf):
                    container.mux(packet)
            for packet in stream.encode():
                container.mux(packet)
        data = sink.getvalue()
        if len(data) > limit:
            raise ValueError("Encoded IMAGE exceeds byte limit")
        return data
    except av.error.FFmpegError:
        raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
    except ValueError:
        raise
    except Exception:
        raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None


def encode_image(value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY,
                 max_bytes=MAX_PAYLOAD_BYTES):
    """Return (descriptor, bytes); reject unsupported inputs, never fall back."""
    fmt, quality = validate_options(transfer_format, transfer_quality)
    if fmt == "exact":
        return codec.encode_payload("IMAGE", value, max_bytes=max_bytes)
    limit = codec._byte_limit(max_bytes)
    if not isinstance(value, torch.Tensor) or value.device.type != "cpu":
        raise ValueError("IMAGE must be a CPU tensor")
    if value.layout != torch.strided or value.is_nested:
        raise ValueError("IMAGE must have dense strided layout")
    if value.dtype not in codec._DTYPE_NAMES:
        raise ValueError("Unsupported IMAGE dtype")
    shape = list(value.shape)
    codec._tensor_size("IMAGE", shape, value.element_size(), limit)
    alpha_size = shape[0] * shape[1] * shape[2] * value.element_size() if shape[-1] == 4 else 0
    itemsize = value.element_size()
    ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
    if not torch.isfinite(value).all().item():
        raise ValueError("IMAGE values must be finite")
    value = value.detach().resolve_neg()
    alpha, alpha_payload = None, b""
    if shape[-1] == 4:
        ad, alpha_payload = codec.encode_payload("IMAGE", value[..., 3:4], max_bytes=limit)
        alpha = {"descriptor": ad, "size": len(alpha_payload)}
    import numpy as np
    def frames():
        for frame in value:
            arr = frame[..., :3].float().numpy()
            with np.errstate(over="ignore", invalid="ignore"):
                yield np.ascontiguousarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8))
    budget = limit - len(alpha_payload)
    if fmt in _VIDEO_CODECS:
        parts = [_encode_video(frames(), fmt, quality, ph, pw, budget)]
        header = _video_header(parts[0], fmt, ph, pw, shape[0])
        coded_frames = _coded_frame_bound(parts[0], fmt, header)
        _shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
    else:
        parts = []
        for frame in frames():
            part = _encode_still(frame, fmt, quality, budget)
            budget -= len(part)
            parts.append(part)
    descriptor = {"encoding": "image-transfer-v1", "format": fmt, "quality": quality,
                  "shape": shape, "layout": "BHWC", "dtype": codec._DTYPE_NAMES[value.dtype],
                  "parts": [len(p) for p in parts], "alpha": alpha}
    _descriptor_budget(descriptor)
    payload = b"".join(parts) + alpha_payload
    if len(payload) > limit:
        raise ValueError("Encoded IMAGE exceeds byte limit")
    return descriptor, payload


def _still_header(data, fmt):
    """Read dimensions before Pillow (WebP may allocate pixels during open)."""
    if fmt in ("png", "legacy_png"):
        if data[:8] != b"\x89PNG\r\n\x1a\n":
            raise ValueError("Invalid IMAGE PNG header")
        offset, dimensions, saw_data, chunks = 8, None, False, 0
        while offset < len(data):
            chunks += 1
            if chunks > MAX_CONTAINER_ELEMENTS:
                raise ValueError("IMAGE PNG chunk count exceeds limit")
            if offset + 12 > len(data):
                raise ValueError("Truncated IMAGE PNG chunk")
            length = int.from_bytes(data[offset:offset + 4], "big")
            end = offset + 12 + length
            if end > len(data):
                raise ValueError("Truncated IMAGE PNG chunk")
            tag = data[offset + 4:offset + 8]
            body = data[offset + 8:end - 4]
            if zlib.crc32(data[offset + 4:end - 4]) != int.from_bytes(data[end - 4:end], "big"):
                raise ValueError("Invalid IMAGE PNG checksum")
            if tag == b"IHDR" and offset == 8 and length == 13:
                width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", body)
                if depth != 8 or color not in (0, 2) or (compression, filtering, interlace) != (0, 0, 0):
                    raise ValueError("Unsupported IMAGE PNG layout")
                dimensions = (height, width)
            elif tag == b"IDAT" and dimensions:
                saw_data = True
            elif tag == b"IEND" and not length and end == len(data) and saw_data:
                return dimensions
            else:
                raise ValueError("Unsupported IMAGE PNG chunk")
            offset = end
        raise ValueError("Truncated IMAGE PNG")
    if fmt == "webp":
        # Our encoder writes a single lossy VP8 still. Reject animation/alpha/
        # extended WebP before WebPAnimDecoder can allocate its canvas.
        if len(data) < 30 or data[:4] != b"RIFF" or data[8:16] != b"WEBPVP8 ":
            raise ValueError("Unsupported IMAGE WebP header")
        size = int.from_bytes(data[16:20], "little")
        if int.from_bytes(data[4:8], "little") + 8 != len(data) or 20 + size + size % 2 != len(data):
            raise ValueError("Invalid IMAGE WebP length")
        if data[23:26] != b"\x9d\x01\x2a":
            raise ValueError("Invalid IMAGE WebP frame")
        width, height = struct.unpack("<HH", data[26:30])
        return height & 0x3fff, width & 0x3fff
    if data[:2] != b"\xff\xd8" or data[-2:] != b"\xff\xd9":
        raise ValueError("Invalid IMAGE JPEG header or trailer")
    offset, dimensions = 2, None
    while offset + 4 <= len(data):
        if data[offset] != 0xff:
            raise ValueError("Invalid IMAGE JPEG marker")
        marker = data[offset + 1]
        size = int.from_bytes(data[offset + 2:offset + 4], "big")
        end = offset + 2 + size
        if size < 2 or end > len(data):
            raise ValueError("Truncated IMAGE JPEG marker")
        if marker in (0xc0, 0xc2):
            if dimensions or size < 8 or data[offset + 4] != 8 or data[offset + 9] not in (1, 3):
                raise ValueError("Unsupported IMAGE JPEG layout")
            dimensions = struct.unpack(">HH", data[offset + 5:offset + 9])
        elif marker == 0xda:
            if dimensions is None:
                raise ValueError("Missing IMAGE JPEG dimensions")
            return dimensions
        offset = end
    raise ValueError("Missing IMAGE JPEG pixels")


def _video_header(data, fmt, height, width, count):
    """Strict bounded EBML preflight: one track, no lacing/attachments/encodings.

    Native stream probing is not the first dimension check: it may instantiate a
    decoder. Validate the Matroska header and count blocks without decompressing.
    """
    nodes = 0
    def vint(pos, end, identifier=False):
        if pos >= end or data[pos] == 0:
            raise ValueError("Invalid IMAGE EBML integer")
        length = 9 - data[pos].bit_length()
        if length > (4 if identifier else 8) or pos + length > end:
            raise ValueError("Truncated IMAGE EBML integer")
        value = int.from_bytes(data[pos:pos + length], "big")
        if not identifier:
            value &= (1 << (7 * length)) - 1
            if value == (1 << (7 * length)) - 1:
                raise ValueError("Unknown-sized IMAGE EBML elements are forbidden")
        return value, pos + length
    def elements(start, end):
        nonlocal nodes
        while start < end:
            nodes += 1
            if nodes > MAX_CONTAINER_ELEMENTS:
                raise ValueError("IMAGE container element count exceeds limit")
            tag, pos = vint(start, end, True)
            length, body = vint(pos, end)
            finish = body + length
            if finish > end:
                raise ValueError("Truncated IMAGE EBML element")
            yield tag, body, finish
            start = finish
    def integer(start, end):
        if not 1 <= end - start <= 8:
            raise ValueError("Invalid IMAGE EBML numeric field")
        return int.from_bytes(data[start:end], "big")
    roots = list(elements(0, len(data)))
    if len(roots) != 2 or [r[0] for r in roots] != [0x1a45dfa3, 0x18538067]:
        raise ValueError("Invalid IMAGE Matroska structure")
    tracks, blocks, packets = [], 0, []
    for tag, start, end in elements(roots[1][1], roots[1][2]):
        if tag == 0x1654ae6b:
            for entry, es, ee in elements(start, end):
                if entry in (0xbf, 0xec):
                    continue
                if entry != 0xae:
                    raise ValueError("Unexpected IMAGE track element")
                fields, dims = {}, {}
                for field, fs, fe in elements(es, ee):
                    if field in fields:
                        raise ValueError("Duplicate IMAGE track field")
                    fields[field] = (fs, fe)
                    if field == 0x6d80:
                        raise ValueError("IMAGE track encodings are forbidden")
                    if field == 0xe0:
                        for dim, ds, de in elements(fs, fe):
                            if dim in dims:
                                raise ValueError("Duplicate IMAGE video field")
                            dims[dim] = (ds, de)
                if any(key not in fields for key in (0xd7, 0x83, 0x86)) or any(key not in dims for key in (0xb0, 0xba)):
                    raise ValueError("Incomplete IMAGE video track")
                if integer(*fields[0xd7]) != 1 or integer(*fields[0x83]) != 1:
                    raise ValueError("Unsupported IMAGE video track")
                if fields[0x86][1] - fields[0x86][0] > 32:
                    raise ValueError("IMAGE codec identifier exceeds limit")
                codec_id = bytes(data[slice(*fields[0x86])])
                if codec_id != _CODEC_IDS[fmt] and not (fmt == "ffv1" and codec_id == b"V_MS/VFW/FOURCC"):
                    raise ValueError("IMAGE video codec mismatch")
                if integer(*dims[0xb0]) != width or integer(*dims[0xba]) != height:
                    raise ValueError("IMAGE video header dimensions mismatch")
                if codec_id == b"V_MS/VFW/FOURCC":
                    cs, ce = fields.get(0x63a2, (0, 0))
                    if ce - cs < 40 or data[cs + 16:cs + 20] != b"FFV1":
                        raise ValueError("Invalid IMAGE FFV1 codec header")
                    if struct.unpack("<ii", data[cs + 4:cs + 12]) != (width, height):
                        raise ValueError("IMAGE FFV1 header dimensions mismatch")
                tracks.append(fields)
        elif tag == 0x1f43b675:
            def block(bs, be):
                nonlocal blocks
                track, pos = vint(bs, be)
                if track != 1 or pos + 3 >= be or data[pos + 2] & 0x06:
                    raise ValueError("Invalid or laced IMAGE video block")
                blocks += 1
                if blocks > count:
                    raise ValueError("IMAGE video frame count exceeds descriptor")
                packets.append((pos + 3, be))
            for field, fs, fe in elements(start, end):
                if field == 0xa3:
                    block(fs, fe)
                elif field == 0xa0:
                    group = [e for e in elements(fs, fe) if e[0] == 0xa1]
                    if len(group) != 1:
                        raise ValueError("Invalid IMAGE video block group")
                    block(group[0][1], group[0][2])
                elif field not in (0xe7, 0xec, 0xbf):
                    raise ValueError("Unsupported IMAGE cluster field")
        elif tag not in (0x114d9b74, 0x1549a966, 0x1c53bb6b, 0x1254c367, 0xec, 0xbf):
            raise ValueError("Unsupported IMAGE Matroska element")
    if len(tracks) != 1 or blocks != count:
        raise ValueError("IMAGE video track or frame count mismatch")
    cs, ce = tracks[0].get(0x63a2, (0, 0))
    if not 0 < ce - cs <= MAX_DESCRIPTOR_BYTES:
        raise ValueError("IMAGE codec configuration exceeds limit or is missing")
    if bytes(data[slice(*tracks[0][0x86])]) == b"V_MS/VFW/FOURCC":
        cs += 40  # Skip BITMAPINFOHEADER; FFV1 receives only its codec config.
    return bytes(data[cs:ce]), packets


def _coded_frame_bound(data, fmt, header):
    """Conservative per-packet allocation bound before native decode.

    FFV1 has one picture per packet. Every AVC/HEVC picture needs at least one
    VCL NAL; every AV1 output (including show-existing) needs a frame/header OBU.
    Counting slices overestimates pictures, which is safe. Do not trust a
    Matroska block count alone: a malicious packet may contain several pictures.
    """
    config, packets = header
    if fmt == "ffv1":
        return len(packets)
    total, units = 0, 0
    if fmt in ("h264", "h265"):
        index = 4 if fmt == "h264" else 21
        if len(config) <= index or config[0] != 1:
            raise ValueError("Invalid IMAGE video configuration")
        length_size = (config[index] & 3) + 1
    for start, end in packets:
        while start < end:
            units += 1
            if units > MAX_CONTAINER_ELEMENTS:
                raise ValueError("IMAGE coded unit count exceeds limit")
            if fmt in ("h264", "h265"):
                if start + length_size > end:
                    raise ValueError("Truncated IMAGE NAL length")
                length = int.from_bytes(data[start:start + length_size], "big")
                start += length_size
                if length < (1 if fmt == "h264" else 2) or length > end - start:
                    raise ValueError("Truncated IMAGE NAL unit")
                kind = (data[start] & 31) if fmt == "h264" else (data[start] >> 1) & 63
                total += (1 <= kind <= 5) if fmt == "h264" else (kind <= 31)
            else:
                tag = data[start]
                start += 1
                if tag & 0x81 or not tag & 2:
                    raise ValueError("Unsupported IMAGE AV1 OBU header")
                kind = (tag >> 3) & 15
                if tag & 4:
                    start += 1
                length = 0
                for shift in range(0, 56, 7):
                    if start >= end:
                        raise ValueError("Truncated IMAGE AV1 OBU length")
                    value = data[start]
                    start += 1
                    length |= (value & 127) << shift
                    if not value & 128:
                        break
                else:
                    raise ValueError("IMAGE AV1 OBU length exceeds limit")
                if length > end - start:
                    raise ValueError("Truncated IMAGE AV1 OBU")
                total += kind in (3, 6)
            start += length
            if total > MAX_FRAMES * 4:
                raise ValueError("IMAGE coded frame count exceeds limit")
    if total < len(packets):
        raise ValueError("IMAGE video has missing coded frames")
    return total


def _decode_video(data, fmt, height, width, count, header):
    av = _import_av()
    maximum = ((height + 63) // 64 * 64) * ((width + 63) // 64 * 64)
    ctx = av.CodecContext.create("libdav1d" if fmt == "av1" else "hevc" if fmt == "h265" else fmt, "r")
    ctx.width, ctx.height = width, height
    ctx.extradata, packets = header
    ctx.thread_count = 1
    ctx.options = {"max_pixels": str(maximum), "threads": "1", "err_detect": "explode"}
    ctx.open()
    if "max_pixels" in ctx.options:
        raise ValueError("IMAGE decoder does not support max_pixels safety limit")
    if (ctx.width, ctx.height) != (width, height):
        raise ValueError("IMAGE bitstream dimensions mismatch")
    decoded = 0
    for bounds in packets + [None]:
        packet = av.Packet(data[slice(*bounds)]) if bounds is not None else None
        for frame in ctx.decode(packet):
            if decoded >= count or (frame.width, frame.height) != (width, height):
                raise ValueError("IMAGE decoded dimensions or frame count mismatch")
            decoded += 1
            yield frame.to_ndarray(format="rgb24")
    if decoded != count:
        raise ValueError("IMAGE decoded frame count mismatch")


def decode_image(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
    """Decode to an independent contiguous CPU tensor after bounded preflight."""
    limit = codec._byte_limit(max_bytes)
    if not isinstance(descriptor, dict):
        raise ValueError("IMAGE descriptor must be a dictionary")
    if descriptor.get("encoding") == "tensor-le-v1":
        return codec.decode_payload("IMAGE", descriptor, payload, max_bytes=limit)
    _descriptor_budget(descriptor)
    if set(descriptor) != {"encoding", "format", "quality", "shape", "layout", "dtype", "parts", "alpha"}:
        raise ValueError("Invalid IMAGE descriptor fields")
    if descriptor["encoding"] != "image-transfer-v1" or descriptor["layout"] != "BHWC":
        raise ValueError("Invalid IMAGE descriptor encoding or layout")
    fmt, _ = validate_options(descriptor["format"], descriptor["quality"])
    if fmt == "exact":
        raise ValueError("Exact IMAGE requires the tensor encoding")
    if not isinstance(payload, bytes):
        raise ValueError("IMAGE payload must be bytes")
    if len(payload) > limit:
        raise ValueError("IMAGE payload exceeds byte limit")
    shape = descriptor["shape"]
    if type(descriptor["dtype"]) is not str or descriptor["dtype"] not in codec._DTYPES:
        raise ValueError("Unsupported IMAGE descriptor dtype")
    dtype, itemsize = codec._DTYPES[descriptor["dtype"]]
    codec._tensor_size("IMAGE", shape, itemsize, limit)
    alpha_size, _ = _alpha_metadata(descriptor["alpha"], shape, limit)
    if alpha_size and descriptor["alpha"]["descriptor"]["dtype"] != descriptor["dtype"]:
        raise ValueError("IMAGE alpha dtype mismatch")
    ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
    count, height, width, channels = shape
    lengths = descriptor["parts"]
    if type(lengths) is not list or len(lengths) != (1 if fmt in _VIDEO_CODECS else count):
        raise ValueError("IMAGE part count mismatch")
    total = alpha_size
    for length in lengths:
        if type(length) is not int or not 0 < length <= limit - total:
            raise ValueError("IMAGE part length exceeds byte limit")
        total += length
    if total != len(payload):
        raise ValueError("IMAGE payload length does not match parts")
    parts, offset = [], 0
    for length in lengths:
        parts.append(memoryview(payload)[offset:offset + length])
        offset += length
    try:
        if fmt in _VIDEO_CODECS:
            video_header = _video_header(parts[0], fmt, ph, pw, count)
            coded_frames = _coded_frame_bound(parts[0], fmt, video_header)
            _shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
        else:
            for part in parts:
                if _still_header(part, fmt) != (height, width):
                    raise ValueError("IMAGE header dimensions mismatch")
        # All headers and aggregate budgets are checked before output allocation.
        alpha = None
        if alpha_size:
            alpha = codec.decode_payload("IMAGE", descriptor["alpha"]["descriptor"], payload[offset:], max_bytes=limit)
        result = torch.empty(shape, dtype=dtype, device="cpu")
        if fmt in _VIDEO_CODECS:
            frames = _decode_video(parts[0], fmt, ph, pw, count, video_header)
        else:
            def stills():
                import numpy as np
                from PIL import Image
                for part in parts:
                    with Image.open(io.BytesIO(part)) as image:
                        if image.size != (width, height) or getattr(image, "n_frames", 1) != 1:
                            raise ValueError("IMAGE still dimensions or frame count mismatch")
                        if image.mode not in ("L", "RGB"):
                            raise ValueError("Unsupported IMAGE still pixel mode")
                        yield np.array(image.convert("RGB"), dtype=np.uint8)
            frames = stills()
        color = 3 if channels == 4 else channels
        for index, frame in enumerate(frames):
            # Float32 division followed by the source dtype cast reproduces
            # Collector normalized pixels without narrowing exact alpha.
            result[index, ..., :color] = torch.from_numpy(frame[:height, :width, :color].copy()).float().div_(255)
        if alpha is not None:
            result[..., 3:4] = alpha
        return result
    except ValueError:
        raise
    except Exception:
        raise ValueError("Malformed or unsupported IMAGE compressed payload") from None