"""Split Annex-B H.264 into access units."""
from __future__ import annotations


def _start_at(buf: bytes, i: int) -> int:
    n = len(buf)
    while i + 2 < n:
        if buf[i] == 0 and buf[i + 1] == 0:
            if buf[i + 2] == 1:
                return i
            if i + 3 < n and buf[i + 2] == 0 and buf[i + 3] == 1:
                return i
        i += 1
    return -1


def nal_unit_type(nal: bytes) -> int:
    """First NAL type in an AU (5 = IDR)."""
    i = _start_at(nal, 0)
    if i < 0:
        return 0
    hdr = i + (4 if i + 3 < len(nal) and nal[i + 2] == 0 else 3)
    if hdr >= len(nal):
        return 0
    return nal[hdr] & 0x1F


def is_keyframe(au: bytes) -> bool:
    i = 0
    n = len(au)
    while True:
        s = _start_at(au, i)
        if s < 0:
            return False
        hdr = s + (4 if s + 3 < n and au[s + 2] == 0 else 3)
        if hdr >= n:
            return False
        ntype = au[hdr] & 0x1F
        if ntype == 5:
            return True
        i = hdr + 1


class AnnexBSplitter:
    """Feed byte-stream; yield complete AUs (including start codes)."""

    def __init__(self) -> None:
        self._buf = bytearray()

    def push(self, data: bytes) -> list[bytes]:
        if data:
            self._buf.extend(data)
        out: list[bytes] = []
        buf = self._buf
        first = _start_at(bytes(buf), 0)
        if first < 0:
            if len(buf) > 1_000_000:
                del buf[: len(buf) - 8]
            return out
        if first:
            del buf[:first]
        pos = 0
        while True:
            nxt = _start_at(bytes(buf), pos + 3)
            if nxt < 0:
                break
            au = bytes(buf[:nxt])
            if au:
                out.append(au)
            del buf[:nxt]
            pos = 0
        return out