// Annex-B helpers for H.264 / H.265 video framing.
//
// The host encodes H.264/H.265 via ffmpeg (hwcodec crate) and the default
// muxer-less output is Annex-B: NAL units separated by 3- or 4-byte start
// codes (`00 00 01` or `00 00 00 01`). SPS/PPS are inlined before each IDR.
//
// WebCodecs' VideoDecoder accepts Annex-B chunks directly when configured
// without a `description` field — in that mode the browser scans for start
// codes itself. We just need to:
//   1. Find the first SPS NAL inside a keyframe to derive the avc1/hvc1
//      codec string the decoder requires.
//   2. Pass each frame's bytes into the decoder unmodified.

const NAL_UNIT_TYPE_MASK_H264 = 0x1f;
const SPS_NAL_TYPE_H264 = 7;

/** Iterate NAL units in an Annex-B stream. Yields `payload` excluding the
 *  start code; first byte is the NAL header. */
export function* iterAnnexBNalu(data: Uint8Array): Generator<Uint8Array> {
  // Locate every start code (00 00 01 or 00 00 00 01), then yield the
  // bytes between consecutive start codes.
  const starts: number[] = [];
  for (let i = 0; i + 2 < data.length; i++) {
    if (data[i] === 0 && data[i + 1] === 0) {
      if (data[i + 2] === 1) {
        starts.push(i + 3);
        i += 2;
      } else if (i + 3 < data.length && data[i + 2] === 0 && data[i + 3] === 1) {
        starts.push(i + 4);
        i += 3;
      }
    }
  }
  for (let k = 0; k < starts.length; k++) {
    const start = starts[k]!;
    const end = k + 1 < starts.length
      ? starts[k + 1]! - (data[starts[k + 1]! - 4] === 0 ? 4 : 3)
      : data.length;
    if (end > start) {
      yield data.subarray(start, end);
    }
  }
}

/**
 * Build an `avc1.PPCCLL` codec string from a keyframe's SPS NAL unit.
 *
 * SPS layout (after NAL header byte):
 *   byte 0: profile_idc
 *   byte 1: constraint_set_flags (low 6 bits) | reserved_zero_2bits
 *   byte 2: level_idc
 *
 * Returns null if no SPS found in the buffer.
 */
export function deriveH264CodecString(keyframe: Uint8Array): string | null {
  for (const nalu of iterAnnexBNalu(keyframe)) {
    if (nalu.length < 4) continue;
    const naluType = nalu[0]! & NAL_UNIT_TYPE_MASK_H264;
    if (naluType !== SPS_NAL_TYPE_H264) continue;
    const profile = nalu[1]!;
    const constraints = nalu[2]!;
    const level = nalu[3]!;
    const hex = (n: number): string => n.toString(16).padStart(2, "0").toUpperCase();
    return `avc1.${hex(profile)}${hex(constraints)}${hex(level)}`;
  }
  return null;
}
