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

using System.IO;
using System.Text;

namespace Vav1Player.Container;

public class MatroskaTrackInfo
{
    public uint TrackNumber { get; set; }
    public string? CodecId { get; set; }
    public uint PixelWidth { get; set; }
    public uint PixelHeight { get; set; }
    public double Duration { get; set; }
    public byte[]? CodecPrivate { get; set; }
    public List<MatroskaBlock> Blocks { get; set; } = new List<MatroskaBlock>();
}

public struct MatroskaBlock
{
    public long Offset { get; set; }
    public int Size { get; set; }
    public ulong Timestamp { get; set; }
    public bool IsKeyFrame { get; set; }
    public byte[] Data { get; set; }
}

public class MatroskaParser
{
    private readonly byte[] _fileData;
    private int _position;

    // EBML Element IDs
    private static readonly Dictionary<uint, string> ElementIds = new Dictionary<uint, string>
    {
        { 0x1A45DFA3, "EBML" },
        { 0x18538067, "Segment" },
        { 0x1549A966, "Info" },
        { 0x1654AE6B, "Tracks" },
        { 0x1F43B675, "Cluster" },
        { 0xAE, "TrackEntry" },
        { 0xD7, "TrackNumber" },
        { 0x83, "TrackType" },
        { 0x86, "CodecID" },
        { 0x63A2, "CodecPrivate" },
        { 0xB0, "PixelWidth" },
        { 0xBA, "PixelHeight" },
        { 0x4489, "Duration" },
        { 0xE7, "Timestamp" },
        { 0xA3, "SimpleBlock" },
        { 0xA1, "Block" },
        { 0xA0, "BlockGroup" }
    };

    public MatroskaParser(byte[] fileData)
    {
        _fileData = fileData;
        _position = 0;
    }

    public List<MatroskaTrackInfo> Parse()
    {
        var tracks = new List<MatroskaTrackInfo>();

        while (_position < _fileData.Length)
        {
            var element = ReadElement();
            if (element.Id == 0x18538067) // Segment
            {
                ParseSegment(element, tracks);
            }
            else
            {
                SkipElement(element);
            }
        }

        return tracks;
    }

    private EbmlElement ReadElement()
    {
        if (_position >= _fileData.Length)
            throw new EndOfStreamException();

        uint id = ReadElementId();
        ulong size = ReadElementSize();
        long dataOffset = _position;

        return new EbmlElement
        {
            Id = id,
            Size = size,
            DataOffset = dataOffset
        };
    }

    private uint ReadElementId()
    {
        if (_position >= _fileData.Length)
            throw new EndOfStreamException();

        byte firstByte = _fileData[_position];
        int idLength = GetElementIdLength(firstByte);

        if (_position + idLength > _fileData.Length)
            throw new EndOfStreamException();

        uint id = 0;
        for (int i = 0; i < idLength; i++)
        {
            id = (id << 8) | _fileData[_position + i];
        }

        _position += idLength;
        return id;
    }

    private ulong ReadElementSize()
    {
        if (_position >= _fileData.Length)
            throw new EndOfStreamException();

        byte firstByte = _fileData[_position];
        int sizeLength = GetElementSizeLength(firstByte);

        if (_position + sizeLength > _fileData.Length)
            throw new EndOfStreamException();

        ulong size = 0;
        byte mask = (byte)(0xFF >> sizeLength);

        size = (ulong)(firstByte & mask);
        for (int i = 1; i < sizeLength; i++)
        {
            size = (size << 8) | _fileData[_position + i];
        }

        _position += sizeLength;
        return size;
    }

    private int GetElementIdLength(byte firstByte)
    {
        if ((firstByte & 0x80) != 0) return 1;
        if ((firstByte & 0x40) != 0) return 2;
        if ((firstByte & 0x20) != 0) return 3;
        if ((firstByte & 0x10) != 0) return 4;
        throw new InvalidDataException("Invalid EBML element ID");
    }

    private int GetElementSizeLength(byte firstByte)
    {
        if ((firstByte & 0x80) != 0) return 1;
        if ((firstByte & 0x40) != 0) return 2;
        if ((firstByte & 0x20) != 0) return 3;
        if ((firstByte & 0x10) != 0) return 4;
        if ((firstByte & 0x08) != 0) return 5;
        if ((firstByte & 0x04) != 0) return 6;
        if ((firstByte & 0x02) != 0) return 7;
        if ((firstByte & 0x01) != 0) return 8;
        throw new InvalidDataException("Invalid EBML element size");
    }

    private void ParseSegment(EbmlElement segment, List<MatroskaTrackInfo> tracks)
    {
        long segmentEnd = segment.DataOffset + (long)segment.Size;
        _position = (int)segment.DataOffset;

        while (_position < segmentEnd && _position < _fileData.Length)
        {
            var element = ReadElement();

            switch (element.Id)
            {
                case 0x1654AE6B: // Tracks
                    ParseTracks(element, tracks);
                    break;
                case 0x1F43B675: // Cluster
                    ParseCluster(element, tracks);
                    break;
                default:
                    SkipElement(element);
                    break;
            }
        }
    }

    private void ParseTracks(EbmlElement tracksElement, List<MatroskaTrackInfo> tracks)
    {
        long tracksEnd = tracksElement.DataOffset + (long)tracksElement.Size;
        _position = (int)tracksElement.DataOffset;

        while (_position < tracksEnd && _position < _fileData.Length)
        {
            var element = ReadElement();

            if (element.Id == 0xAE) // TrackEntry
            {
                var track = ParseTrackEntry(element);
                if (track != null && track.CodecId == "V_AV1")
                {
                    tracks.Add(track);
                }
            }
            else
            {
                SkipElement(element);
            }
        }
    }

    private MatroskaTrackInfo? ParseTrackEntry(EbmlElement trackEntry)
    {
        var track = new MatroskaTrackInfo();
        long trackEnd = trackEntry.DataOffset + (long)trackEntry.Size;
        _position = (int)trackEntry.DataOffset;

        while (_position < trackEnd && _position < _fileData.Length)
        {
            var element = ReadElement();

            switch (element.Id)
            {
                case 0xD7: // TrackNumber
                    track.TrackNumber = ReadUInt(element);
                    break;
                case 0x83: // TrackType
                    uint trackType = ReadUInt(element);
                    // 1 = video, 2 = audio, 3 = complex, 0x10 = logo, 0x11 = subtitle, 0x12 = buttons, 0x20 = control
                    if (trackType != 1) return null; // Only video tracks
                    break;
                case 0x86: // CodecID
                    track.CodecId = ReadString(element);
                    break;
                case 0x63A2: // CodecPrivate
                    track.CodecPrivate = ReadBytes(element);
                    break;
                case 0xB0: // PixelWidth
                    track.PixelWidth = ReadUInt(element);
                    break;
                case 0xBA: // PixelHeight
                    track.PixelHeight = ReadUInt(element);
                    break;
                default:
                    SkipElement(element);
                    break;
            }
        }

        return track.CodecId == "V_AV1" ? track : null;
    }

    private void ParseCluster(EbmlElement cluster, List<MatroskaTrackInfo> tracks)
    {
        long clusterEnd = cluster.DataOffset + (long)cluster.Size;
        _position = (int)cluster.DataOffset;
        ulong clusterTimestamp = 0;

        while (_position < clusterEnd && _position < _fileData.Length)
        {
            var element = ReadElement();

            switch (element.Id)
            {
                case 0xE7: // Timestamp
                    clusterTimestamp = ReadULong(element);
                    break;
                case 0xA3: // SimpleBlock
                    ParseSimpleBlock(element, tracks, clusterTimestamp);
                    break;
                case 0xA0: // BlockGroup
                    ParseBlockGroup(element, tracks, clusterTimestamp);
                    break;
                default:
                    SkipElement(element);
                    break;
            }
        }
    }

    private void ParseSimpleBlock(EbmlElement blockElement, List<MatroskaTrackInfo> tracks, ulong clusterTimestamp)
    {
        long blockOffset = blockElement.DataOffset;
        int blockSize = (int)blockElement.Size;

        if (blockSize < 4) return;

        _position = (int)blockOffset;

        // Read track number
        uint trackNumber = (uint)ReadVInt();

        // Read timestamp (relative to cluster timestamp)
        short relativeTimestamp = (short)((_fileData[_position] << 8) | _fileData[_position + 1]);
        _position += 2;

        // Read flags
        byte flags = _fileData[_position];
        _position++;

        bool isKeyFrame = (flags & 0x80) != 0;

        // Find the track
        var track = tracks.FirstOrDefault(t => t.TrackNumber == trackNumber);
        if (track == null) return;

        // Read frame data
        int frameDataSize = blockSize - (_position - (int)blockOffset);
        if (frameDataSize <= 0) return;

        byte[] frameData = new byte[frameDataSize];
        Array.Copy(_fileData, _position, frameData, 0, frameDataSize);

        var block = new MatroskaBlock
        {
            Offset = _position,
            Size = frameDataSize,
            Timestamp = clusterTimestamp + (ulong)relativeTimestamp,
            IsKeyFrame = isKeyFrame,
            Data = frameData
        };

        track.Blocks.Add(block);
        SkipElement(blockElement);
    }

    private void ParseBlockGroup(EbmlElement blockGroup, List<MatroskaTrackInfo> tracks, ulong clusterTimestamp)
    {
        long blockGroupEnd = blockGroup.DataOffset + (long)blockGroup.Size;
        _position = (int)blockGroup.DataOffset;

        while (_position < blockGroupEnd && _position < _fileData.Length)
        {
            var element = ReadElement();

            if (element.Id == 0xA1) // Block
            {
                ParseBlock(element, tracks, clusterTimestamp, true); // BlockGroup blocks are typically keyframes
            }
            else
            {
                SkipElement(element);
            }
        }
    }

    private void ParseBlock(EbmlElement blockElement, List<MatroskaTrackInfo> tracks, ulong clusterTimestamp, bool isKeyFrame)
    {
        long blockOffset = blockElement.DataOffset;
        int blockSize = (int)blockElement.Size;

        if (blockSize < 4) return;

        _position = (int)blockOffset;

        // Read track number
        uint trackNumber = (uint)ReadVInt();

        // Read timestamp
        short relativeTimestamp = (short)((_fileData[_position] << 8) | _fileData[_position + 1]);
        _position += 2;

        // Skip flags
        _position++;

        // Find the track
        var track = tracks.FirstOrDefault(t => t.TrackNumber == trackNumber);
        if (track == null) return;

        // Read frame data
        int frameDataSize = blockSize - (_position - (int)blockOffset);
        if (frameDataSize <= 0) return;

        byte[] frameData = new byte[frameDataSize];
        Array.Copy(_fileData, _position, frameData, 0, frameDataSize);

        var block = new MatroskaBlock
        {
            Offset = _position,
            Size = frameDataSize,
            Timestamp = clusterTimestamp + (ulong)relativeTimestamp,
            IsKeyFrame = isKeyFrame,
            Data = frameData
        };

        track.Blocks.Add(block);
        SkipElement(blockElement);
    }

    private ulong ReadVInt()
    {
        if (_position >= _fileData.Length)
            throw new EndOfStreamException();

        byte firstByte = _fileData[_position];
        int length = GetElementSizeLength(firstByte);

        ulong value = 0;
        byte mask = (byte)(0xFF >> length);

        value = (ulong)(firstByte & mask);
        for (int i = 1; i < length; i++)
        {
            if (_position + i >= _fileData.Length)
                throw new EndOfStreamException();
            value = (value << 8) | _fileData[_position + i];
        }

        _position += length;
        return value;
    }

    private uint ReadUInt(EbmlElement element)
    {
        if (element.Size > 4) return 0;

        uint value = 0;
        long elementEnd = element.DataOffset + (long)element.Size;

        for (long i = element.DataOffset; i < elementEnd && i < _fileData.Length; i++)
        {
            value = (value << 8) | _fileData[i];
        }

        SkipElement(element);
        return value;
    }

    private ulong ReadULong(EbmlElement element)
    {
        if (element.Size > 8) return 0;

        ulong value = 0;
        long elementEnd = element.DataOffset + (long)element.Size;

        for (long i = element.DataOffset; i < elementEnd && i < _fileData.Length; i++)
        {
            value = (value << 8) | _fileData[i];
        }

        SkipElement(element);
        return value;
    }

    private string ReadString(EbmlElement element)
    {
        if (element.Size == 0) return string.Empty;

        long elementEnd = Math.Min(element.DataOffset + (long)element.Size, _fileData.Length);
        int length = (int)(elementEnd - element.DataOffset);

        if (length <= 0) return string.Empty;

        string value = Encoding.UTF8.GetString(_fileData, (int)element.DataOffset, length);
        SkipElement(element);
        return value;
    }

    private byte[] ReadBytes(EbmlElement element)
    {
        if (element.Size == 0) return Array.Empty<byte>();

        long elementEnd = Math.Min(element.DataOffset + (long)element.Size, _fileData.Length);
        int length = (int)(elementEnd - element.DataOffset);

        if (length <= 0) return Array.Empty<byte>();

        byte[] value = new byte[length];
        Array.Copy(_fileData, (int)element.DataOffset, value, 0, length);
        SkipElement(element);
        return value;
    }

    private void SkipElement(EbmlElement element)
    {
        _position = (int)(element.DataOffset + (long)element.Size);
    }

    public byte[] GetBlockData(MatroskaBlock block)
    {
        return block.Data;
    }
}

internal struct EbmlElement
{
    public uint Id { get; set; }
    public ulong Size { get; set; }
    public long DataOffset { get; set; }
}