// Package h264 parses Annex-B H.264 byte streams (as produced by libx264
// with annexb=1) and groups the raw NAL units into access units (AUs), where
// each AU corresponds to one decoded video frame. This matches the packet
// boundaries the reference driver receives from FFmpeg's demuxer.
//
// AU grouping rules:
//
//   - A VCL NAL (slice: types 1-5) begins a new AU.
//   - Non-VCL NALs (SPS/PPS/SEI, ...) are attached as a prefix to the AU
//     that contains the next VCL NAL, so a keyframe AU naturally becomes
//     SPS + PPS + IDR slice.
//   - An access unit delimiter (type 9) is an explicit boundary.
package h264

// NAL unit types relevant for AU assembly.
const (
	nalTypeSlice    = 1
	nalTypeIDRSlice = 5
	nalTypeSEI      = 6
	nalTypeSPS      = 7
	nalTypePPS      = 8
	nalTypeAUD      = 9
)

// isVCL reports whether a NAL unit type carries slice data.
func isVCL(t byte) bool {
	return t >= nalTypeSlice && t <= nalTypeIDRSlice
}

// nalType extracts the NAL unit type from the first payload byte.
func nalType(payloadStart []byte) byte {
	if len(payloadStart) == 0 {
		return 0
	}
	return payloadStart[0] & 0x1F
}

// payloadStarts returns the index of the byte immediately following each
// start code (00 00 01, optionally prefixed with an extra 00).
func payloadStarts(data []byte) []int {
	var starts []int
	for i := 0; i+2 < len(data); i++ {
		if data[i] == 0 && data[i+1] == 0 && data[i+2] == 1 {
			starts = append(starts, i+3)
			i = i + 2
		}
	}
	return starts
}

// startCodeLen determines the length of the start code that immediately
// precedes the payload at index pay in data.
func startCodeLen(data []byte, pay int) int {
	if pay >= 4 && data[pay-4] == 0 && data[pay-3] == 0 && data[pay-2] == 0 && data[pay-1] == 1 {
		return 4
	}
	return 3
}

// SplitAUs splits an Annex-B stream into access units. The returned slices
// alias the input buffer, so callers must not modify data while they are in
// use.
func SplitAUs(data []byte) [][]byte {
	starts := payloadStarts(data)
	if len(starts) == 0 {
		if len(data) == 0 {
			return nil
		}
		return [][]byte{data}
	}

	// Pre-compute the byte index where each NAL (including its start code)
	// begins.
	codeStart := make([]int, len(starts))
	for i, pay := range starts {
		codeStart[i] = pay - startCodeLen(data, pay)
	}

	var aus [][]byte
	auStart := codeStart[0]
	hasVCL := false

	// flush ends the current AU at the start of NAL i and begins a new one.
	flush := func(i int) {
		end := codeStart[i]
		if end > auStart {
			aus = append(aus, data[auStart:end])
		}
		auStart = end
		hasVCL = false
	}

	for i, pay := range starts {
		switch t := nalType(data[pay:]); {
		case t == nalTypeAUD:
			// An access unit delimiter is an explicit boundary; it becomes
			// the prefix of the AU that follows it.
			if hasVCL {
				flush(i)
			}
		case isVCL(t):
			// A slice always starts a new AU once the current one already
			// contains slice data.
			if hasVCL {
				flush(i)
			}
			hasVCL = true
		default:
			// Non-VCL prefix: SPS, PPS, SEI, ... attaches to the current AU.
		}
	}

	if len(data) > auStart {
		aus = append(aus, data[auStart:])
	}
	return aus
}