"""Split AV1 low-overhead OBU streams into temporal units."""
from __future__ import annotations

OBU_SEQUENCE_HEADER = 1
OBU_TEMPORAL_DELIMITER = 2


def _leb128(data: bytes, i: int) -> tuple[int, int] | None:
    n = len(data)
    result = 0
    shift = 0
    while i < n:
        byte = data[i]
        i += 1
        result |= (byte & 0x7F) << shift
        if (byte & 0x80) == 0:
            return result, i
        shift += 7
        if shift > 56:
            return None
    return None


def parse_obus(data: bytes) -> list[tuple[int, int, int]]:
    """Return (type, start, end) for each complete OBU. Empty if malformed."""
    out: list[tuple[int, int, int]] = []
    i = 0
    n = len(data)
    while i < n:
        start = i
        hdr = data[i]
        i += 1
        if hdr & 0x80:
            return []
        obu_type = (hdr >> 3) & 0x0F
        has_ext = (hdr & 0x04) != 0
        has_size = (hdr & 0x02) != 0
        if has_ext:
            if i >= n:
                return []
            i += 1
        if has_size:
            leb = _leb128(data, i)
            if leb is None:
                return []
            size, i = leb
            end = i + size
            if end > n:
                return []
            i = end
        else:
            i = n
        out.append((obu_type, start, i))
    return out


def is_keyframe(tu: bytes) -> bool:
    return any(t == OBU_SEQUENCE_HEADER for t, _s, _e in parse_obus(tu))


class Av1TuSplitter:
    """Feed obu-stream bytes; yield complete temporal units."""

    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
        while True:
            n = len(buf)
            if n < 1:
                break
            i = 0
            tu_end = -1
            saw_coded = False
            ok = True
            while i < n:
                start = i
                hdr = buf[i]
                i += 1
                if hdr & 0x80:
                    ok = False
                    break
                obu_type = (hdr >> 3) & 0x0F
                has_ext = (hdr & 0x04) != 0
                has_size = (hdr & 0x02) != 0
                if has_ext:
                    if i >= n:
                        ok = False
                        break
                    i += 1
                if has_size:
                    leb = _leb128(bytes(buf), i)
                    if leb is None:
                        ok = False
                        break
                    size, i = leb
                    if i + size > n:
                        ok = False
                        break
                    i = i + size
                else:
                    ok = False
                    break
                if obu_type == OBU_TEMPORAL_DELIMITER and saw_coded:
                    tu_end = start
                    break
                if obu_type != OBU_TEMPORAL_DELIMITER:
                    saw_coded = True
            if not ok:
                if n > 2_000_000:
                    del buf[: n - 16]
                break
            if tu_end < 0:
                break
            tu = bytes(buf[:tu_end])
            if tu:
                out.append(tu)
            del buf[:tu_end]
        return out