video-v1/vav1/Vav1Player/Container/Av1BitstreamParser.cs

namespace Vav1Player.Container;

public static class Av1BitstreamParser
{
    public static List<byte[]> ParseMp4Sample(byte[] sampleData)
    {
        var obuList = new List<byte[]>();
        int offset = 0;

        System.Diagnostics.Debug.WriteLine($"[AV1_PARSER] Starting MP4 sample parse, total size: {sampleData.Length}");

        try
        {
            while (offset < sampleData.Length)
            {
                // Check if remaining data is too small for a valid OBU (minimum 1 byte for LEB128 length + 1 byte header)
                if (sampleData.Length - offset < 2)
                {
                    break;
                }

                // Read OBU length (variable length encoding) - this is the total OBU size including header
                var (obuLength, lengthBytes) = ReadLEB128(sampleData, offset);

                // Validate LEB128 reading results
                if (lengthBytes == 0)
                {
                    break;
                }

                // Check if we have enough data for the length field itself
                if (offset + lengthBytes > sampleData.Length)
                {
                    break;
                }

                // Zero length is only valid for padding OBUs or end markers
                if (obuLength == 0)
                {
                    break;
                }

                // Check if the OBU data itself fits within the sample
                long remainingBytes = sampleData.Length - offset - lengthBytes;
                if (obuLength > remainingBytes)
                {
                    break;
                }

                offset += lengthBytes;

                // Extract the entire OBU (including header)
                if (obuLength > int.MaxValue)
                {
                    break;
                }

                int obuSize = (int)obuLength;
                if (offset + obuSize > sampleData.Length)
                {
                    break;
                }

                byte[] obuData = new byte[obuSize];
                Array.Copy(sampleData, offset, obuData, 0, obuSize);
                obuList.Add(obuData);

                offset += obuSize;
            }
        }
        catch (Exception ex)
        {
            System.Diagnostics.Debug.WriteLine($"[AV1_PARSER] Exception: {ex.Message}");
        }

        return obuList;
    }

    private static (uint length, int bytesRead) ReadLEB128(byte[] data, int offset)
    {
        uint result = 0;
        int shift = 0;
        int bytesRead = 0;

        try
        {
            // Check if we have any data left to read
            if (offset >= data.Length)
            {
                return (0, 0);
            }

            while (offset + bytesRead < data.Length && bytesRead < 8) // Max 8 bytes for LEB128
            {
                byte b = data[offset + bytesRead];
                bytesRead++;

                // Extract the 7 data bits (mask out continuation bit)
                uint bits = (uint)(b & 0x7F);

                // Check for potential overflow before shifting
                if (shift >= 32) // Allow up to 32 bits for uint
                {
                    return (0, 0);
                }

                // Check if adding these bits would overflow
                if (shift == 28 && bits > 0xF) // Only 4 bits left in a 32-bit uint
                {
                    return (0, 0);
                }

                result |= bits << shift;

                // If high bit is not set, this is the last byte
                if ((b & 0x80) == 0)
                    break;

                shift += 7;
            }

            // If we reached the end of data but the last byte had the continuation bit set,
            // this indicates truncated/invalid LEB128 data
            if (bytesRead > 0 && offset + bytesRead - 1 < data.Length)
            {
                byte lastByte = data[offset + bytesRead - 1];
                if ((lastByte & 0x80) != 0)
                {
                    return (0, 0);
                }
            }

            // Sanity check the result - OBU lengths should be reasonable
            if (result > 10 * 1024 * 1024) // 10MB max per OBU seems reasonable
            {
                return (0, 0);
            }

            // Additional check for invalid very small values that might indicate parsing errors
            if (bytesRead > 1 && result == 0)
            {
                return (0, 0);
            }
        }
        catch (Exception ex)
        {
            System.Diagnostics.Debug.WriteLine($"[AV1_PARSER] LEB128 Exception: {ex.Message}");
            return (0, 0);
        }

        return (result, bytesRead);
    }

    public static byte[] CombineOBUs(List<byte[]> obuList)
    {
        if (obuList.Count == 0)
            return Array.Empty<byte>();

        if (obuList.Count == 1)
        {
            // For single OBU from MP4, use as-is (already properly formatted)
            return obuList[0];
        }

        // For multiple OBUs, concatenate them directly
        var combinedList = new List<byte>();

        foreach (var obu in obuList)
        {
            // MP4 OBUs are already properly formatted, just concatenate
            combinedList.AddRange(obu);
        }

        return combinedList.ToArray();
    }

    private static byte[] EnsureOBUHasSizeField(byte[] obuData)
    {
        if (obuData.Length == 0)
            return obuData;

        byte header = obuData[0];
        bool hasSizeField = (header & 0x02) != 0;

        // If OBU already has size field, return as-is
        if (hasSizeField)
            return obuData;

        // Add size field to OBU header
        byte newHeader = (byte)(header | 0x02); // Set has_size_field bit
        var payloadSize = obuData.Length - 1; // Exclude original header
        var sizeBytes = EncodeLEB128((uint)payloadSize);

        var result = new byte[1 + sizeBytes.Length + payloadSize];
        result[0] = newHeader;
        Array.Copy(sizeBytes, 0, result, 1, sizeBytes.Length);
        Array.Copy(obuData, 1, result, 1 + sizeBytes.Length, payloadSize);

        return result;
    }

    private static byte[] EncodeLEB128(uint value)
    {
        var bytes = new List<byte>();

        while (value >= 0x80)
        {
            bytes.Add((byte)((value & 0x7F) | 0x80));
            value >>= 7;
        }
        bytes.Add((byte)(value & 0x7F));

        return bytes.ToArray();
    }

    public static void LogOBUInfo(byte[] obuData, string prefix = "")
    {
        if (obuData.Length == 0)
        {
            System.Diagnostics.Debug.WriteLine($"{prefix}OBU: Empty");
            return;
        }

        byte header = obuData[0];
        int obuType = (header >> 3) & 0xF;
        bool extensionFlag = (header & 0x4) != 0;
        bool hasSizeField = (header & 0x2) != 0;

        string obuTypeName = obuType switch
        {
            0 => "Reserved",
            1 => "Sequence Header",
            2 => "Temporal Delimiter",
            3 => "Frame Header",
            4 => "Tile Group",
            5 => "Metadata",
            6 => "Frame",
            7 => "Redundant Frame Header",
            8 => "Tile List",
            15 => "Padding",
            _ => $"Unknown({obuType})"
        };

        // Show first few bytes for debugging
        var hexData = string.Join(" ", obuData.Take(Math.Min(8, obuData.Length)).Select(b => b.ToString("X2")));

        System.Diagnostics.Debug.WriteLine($"{prefix}OBU: Type={obuTypeName}, Size={obuData.Length}, Extension={extensionFlag}, HasSize={hasSizeField}, Header=0x{header:X2}, FirstBytes=[{hexData}]");
    }
}