import Foundation

/// Converts between Annex-B (start-code delimited) and AVCC (4-byte
/// big-endian length prefixed) H.264 representations.
///
/// VideoToolbox consumes AVCC: each NAL unit is preceded by a 32-bit length.
/// Our depacketizer emits Annex-B (the project's canonical 3-byte start codes),
/// so the decoder feeds AVCC derived here.
enum AvccConverter {
    /// Splits an Annex-B access unit into its NAL units (start codes removed).
    static func nalUnits(_ annexB: [UInt8]) -> [[UInt8]] {
        guard annexB.count >= 4 else { return [] }
        let n = annexB.count

        // Locate every start code (3-byte `00 00 01` and 4-byte `00 00 00 01`).
        var offsets: [Int] = []
        var i = 0
        while i + 2 < n {
            if annexB[i] == 0 && annexB[i + 1] == 0 && annexB[i + 2] == 1 {
                offsets.append(i)
                i += 3
                continue
            }
            if i + 3 < n, annexB[i + 2] == 0, annexB[i + 3] == 1 {
                offsets.append(i)
                i += 4
                continue
            }
            i += 1
        }
        guard !offsets.isEmpty else { return [] }

        var units: [[UInt8]] = []
        for (idx, offset) in offsets.enumerated() {
            let nalStart = offset + 3
            let nalEnd = idx + 1 < offsets.count ? offsets[idx + 1] : n
            let nal = Array(annexB[nalStart..<nalEnd])
            if !nal.isEmpty { units.append(nal) }
        }
        return units
    }

    /// Returns the AVCC form of an Annex-B access unit, or nil if it has no NALs.
    static func toAvcc(_ annexB: [UInt8]) -> [UInt8]? {
        let units = nalUnits(annexB)
        guard !units.isEmpty else { return nil }
        var out: [UInt8] = []
        out.reserveCapacity(annexB.count + units.count * 4)
        for nal in units {
            let len = nal.count
            out.append(UInt8((len >> 24) & 0xFF))
            out.append(UInt8((len >> 16) & 0xFF))
            out.append(UInt8((len >> 8) & 0xFF))
            out.append(UInt8(len & 0xFF))
            out.append(contentsOf: nal)
        }
        return out
    }

    /// Splits an AVCC byte array (4-byte big-endian length prefixes) into NAL units.
    static func fromAvcc(_ avcc: [UInt8]) -> [[UInt8]] {
        var units: [[UInt8]] = []
        var i = 0
        let n = avcc.count
        while i + 4 <= n {
            let len = (Int(avcc[i]) << 24) | (Int(avcc[i + 1]) << 16) | (Int(avcc[i + 2]) << 8) | Int(avcc[i + 3])
            guard len > 0, i + 4 + len <= n else { break }
            units.append(Array(avcc[(i + 4)..<(i + 4 + len)]))
            i += 4 + len
        }
        return units
    }
}