// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Buffers.Binary;
using System.Text;

#nullable enable

namespace Ignitor
{
    public static class RenderBatchReader
    {
        private const int ReferenceFrameSize = 20;

        public static RenderBatch Read(ReadOnlySpan<byte> data)
        {
            var sections = Sections.Parse(data);
            var strings = ReadStringTable(data, sections.GetStringTableIndexes(data));
            var diffs = ReadUpdatedComponents(data, sections.GetUpdatedComponentIndexes(data), strings);
            var frames = ReadReferenceFrames(sections.GetReferenceFrameData(data), strings);
            var disposedComponentIds = ReadDisposedComponentIds(data);
            var disposedEventHandlerIds = ReadDisposedEventHandlerIds(data);
            return new RenderBatch(diffs, frames, disposedComponentIds, disposedEventHandlerIds);
        }

        private static string[] ReadStringTable(ReadOnlySpan<byte> data, ReadOnlySpan<byte> indexes)
        {
            var result = new string[indexes.Length / 4];

            for (var i = 0; i < indexes.Length; i += 4)
            {
                var index = BinaryPrimitives.ReadInt32LittleEndian(indexes.Slice(i, 4));

                // The string table entries are all length-prefixed UTF8 blobs
                var length = (int)ReadUnsignedLEB128(data, index, out var numLEB128Bytes);
                var value = Encoding.UTF8.GetString(data.Slice(index + numLEB128Bytes, length));
                result[i / 4] = value;
            }

            return result;
        }

        private static ArrayRange<RenderTreeDiff> ReadUpdatedComponents(ReadOnlySpan<byte> data, ReadOnlySpan<byte> indexes, string[] strings)
        {
            var result = new RenderTreeDiff[indexes.Length / 4];

            for (var i = 0; i < indexes.Length; i += 4)
            {
                var index = BinaryPrimitives.ReadInt32LittleEndian(indexes.Slice(i, 4));

                var componentId = BinaryPrimitives.ReadInt32LittleEndian(data.Slice(index, 4));
                var editCount = BinaryPrimitives.ReadInt32LittleEndian(data.Slice(index + 4, 4));

                var editData = data.Slice(index + 8);
                var edits = new RenderTreeEdit[editCount];
                for (var j = 0; j < editCount; j++)
                {
                    var type = (RenderTreeEditType)BinaryPrimitives.ReadInt32LittleEndian(editData.Slice(0, 4));
                    var siblingIndex = BinaryPrimitives.ReadInt32LittleEndian(editData.Slice(4, 4));

                    // ReferenceFrameIndex and MoveToSiblingIndex share a slot, so this reads
                    // whichever one applies to the edit type
                    var referenceFrameIndex = BinaryPrimitives.ReadInt32LittleEndian(editData.Slice(8, 4));
                    var removedAttributeName = ReadString(editData.Slice(12, 4), strings);

                    editData = editData.Slice(16);

                    switch (type)
                    {
                        case RenderTreeEditType.UpdateText:
                            edits[j] = RenderTreeEdit.UpdateText(siblingIndex, referenceFrameIndex);
                            break;

                        case RenderTreeEditType.UpdateMarkup:
                            edits[j] = RenderTreeEdit.UpdateMarkup(siblingIndex, referenceFrameIndex);
                            break;

                        case RenderTreeEditType.SetAttribute:
                            edits[j] = RenderTreeEdit.SetAttribute(siblingIndex, referenceFrameIndex);
                            break;

                        case RenderTreeEditType.RemoveAttribute:
                            edits[j] = RenderTreeEdit.RemoveAttribute(siblingIndex, removedAttributeName!);
                            break;

                        case RenderTreeEditType.PrependFrame:
                            edits[j] = RenderTreeEdit.PrependFrame(siblingIndex, referenceFrameIndex);
                            break;

                        case RenderTreeEditType.RemoveFrame:
                            edits[j] = RenderTreeEdit.RemoveFrame(siblingIndex);
                            break;

                        case RenderTreeEditType.StepIn:
                            edits[j] = RenderTreeEdit.StepIn(siblingIndex);
                            break;

                        case RenderTreeEditType.StepOut:
                            edits[j] = RenderTreeEdit.StepOut();
                            break;

                        case RenderTreeEditType.PermutationListEntry:
                            edits[j] = RenderTreeEdit.PermutationListEntry(siblingIndex, referenceFrameIndex);
                            break;

                        case RenderTreeEditType.PermutationListEnd:
                            edits[j] = RenderTreeEdit.PermutationListEnd();
                            break;

                        default:
                            throw new InvalidOperationException("Unknown edit type:" + type);
                    }
                }

                result[i / 4] = new RenderTreeDiff(componentId, ToArrayBuilderSegment(edits));
            }

            return new ArrayRange<RenderTreeDiff>(result, result.Length);
        }

        private static ArrayBuilderSegment<T> ToArrayBuilderSegment<T>(T[] entries)
        {
            var builder = new ArrayBuilder<T>();
            builder.Append(entries, 0, entries.Length);
            return builder.ToSegment(0, entries.Length);
        }

        private static ArrayRange<RenderTreeFrame> ReadReferenceFrames(ReadOnlySpan<byte> data, string[] strings)
        {
            var result = new RenderTreeFrame[data.Length / ReferenceFrameSize];

            for (var i = 0; i < data.Length; i += ReferenceFrameSize)
            {
                var frameData = data.Slice(i, ReferenceFrameSize);

                var type = (RenderTreeFrameType)BinaryPrimitives.ReadInt32LittleEndian(frameData.Slice(0, 4));

                // We want each frame to take up the same number of bytes, so that the
                // recipient can index into the array directly instead of having to
                // walk through it.
                // Since we can fit every frame type into 16 bytes, use that as the
                // common size. For smaller frames, we add padding to expand it to
                // 16 bytes.
                switch (type)
                {
                    case RenderTreeFrameType.Attribute:
                        var attributeName = ReadString(frameData.Slice(4, 4), strings);
                        var attributeValue = ReadString(frameData.Slice(8, 4), strings);
                        var attributeEventHandlerId = BinaryPrimitives.ReadUInt64LittleEndian(frameData.Slice(12, 8));
                        result[i / ReferenceFrameSize] = RenderTreeFrame.Attribute(0, attributeName, attributeValue).WithAttributeEventHandlerId(attributeEventHandlerId);
                        break;

                    case RenderTreeFrameType.Component:
                        var componentSubtreeLength = BinaryPrimitives.ReadInt32LittleEndian(frameData.Slice(4, 4));
                        var componentId = BinaryPrimitives.ReadInt32LittleEndian(frameData.Slice(8, 4)); // Nowhere to put this without creating a ComponentState
                        result[i / ReferenceFrameSize] = RenderTreeFrame.ChildComponent(0, componentType: null)
                            .WithComponentSubtreeLength(componentSubtreeLength)
                            .WithComponent(new ComponentState(componentId));
                        break;

                    case RenderTreeFrameType.ComponentReferenceCapture:
                        // Client doesn't process these, skip.
                        result[i / ReferenceFrameSize] = RenderTreeFrame.ComponentReferenceCapture(0, null, 0);
                        break;

                    case RenderTreeFrameType.Element:
                        var elementSubtreeLength = BinaryPrimitives.ReadInt32LittleEndian(frameData.Slice(4, 4));
                        var elementName = ReadString(frameData.Slice(8, 4), strings);
                        result[i / ReferenceFrameSize] = RenderTreeFrame.Element(0, elementName).WithElementSubtreeLength(elementSubtreeLength);
                        break;

                    case RenderTreeFrameType.ElementReferenceCapture:
                        var referenceCaptureId = ReadString(frameData.Slice(4, 4), strings);
                        result[i / ReferenceFrameSize] = RenderTreeFrame.ElementReferenceCapture(0, null)
                            .WithElementReferenceCaptureId(referenceCaptureId);
                        break;

                    case RenderTreeFrameType.Region:
                        var regionSubtreeLength = BinaryPrimitives.ReadInt32LittleEndian(frameData.Slice(4, 4));
                        result[i / ReferenceFrameSize] = RenderTreeFrame.Region(0).WithRegionSubtreeLength(regionSubtreeLength);
                        break;

                    case RenderTreeFrameType.Text:
                        var text = ReadString(frameData.Slice(4, 4), strings);
                        result[i / ReferenceFrameSize] = RenderTreeFrame.Text(0, text);
                        break;

                    case RenderTreeFrameType.Markup:
                        var markup = ReadString(frameData.Slice(4, 4), strings);
                        result[i / ReferenceFrameSize] = RenderTreeFrame.Markup(0, markup);
                        break;

                    default:
                        throw new ArgumentException($"Unsupported frame type: {type}");
                }
            }

            return new ArrayRange<RenderTreeFrame>(result, result.Length);
        }

        private static ArrayRange<int> ReadDisposedComponentIds(ReadOnlySpan<byte> data)
        {
            return new ArrayRange<int>(Array.Empty<int>(), 0);
        }

        private static ArrayRange<ulong> ReadDisposedEventHandlerIds(ReadOnlySpan<byte> data)
        {
            return new ArrayRange<ulong>(Array.Empty<ulong>(), 0);
        }

        private static string? ReadString(ReadOnlySpan<byte> data, string[] strings)
        {
            var index = BinaryPrimitives.ReadInt32LittleEndian(data.Slice(0, 4));
            return index >= 0 ? strings[index] : null;
        }

        private static uint ReadUnsignedLEB128(ReadOnlySpan<byte> data, int startOffset, out int numBytesRead)
        {
            var result = (uint)0;
            var shift = 0;
            var currentByte = (byte)128;
            numBytesRead = 0;

            for (var count = 0; count < 4 && currentByte >= 128; count++)
            {
                currentByte = data[startOffset + count];
                result += (uint)(currentByte & 0x7f) << shift;
                shift += 7;
                numBytesRead++;
            }

            return result;
        }

        private readonly struct Sections
        {
            public static Sections Parse(ReadOnlySpan<byte> data)
            {
                return new Sections(
                    BinaryPrimitives.ReadInt32LittleEndian(data.Slice(data.Length - 20, 4)),
                    BinaryPrimitives.ReadInt32LittleEndian(data.Slice(data.Length - 16, 4)),
                    BinaryPrimitives.ReadInt32LittleEndian(data.Slice(data.Length - 12, 4)),
                    BinaryPrimitives.ReadInt32LittleEndian(data.Slice(data.Length - 8, 4)),
                    BinaryPrimitives.ReadInt32LittleEndian(data.Slice(data.Length - 4, 4)));
            }

            private readonly int _updatedComponents;
            private readonly int _referenceFrames;
            private readonly int _disposedComponentIds;
            private readonly int _disposedEventHandlerIds;
            private readonly int _strings;

            public Sections(int updatedComponents, int referenceFrames, int disposedComponentIds, int disposedEventHandlerIds, int strings)
            {
                _updatedComponents = updatedComponents;
                _referenceFrames = referenceFrames;
                _disposedComponentIds = disposedComponentIds;
                _disposedEventHandlerIds = disposedEventHandlerIds;
                _strings = strings;
            }

            public ReadOnlySpan<byte> GetUpdatedComponentIndexes(ReadOnlySpan<byte> data)
            {
                // This is count-prefixed contiguous array of of integers.
                var count = BinaryPrimitives.ReadInt32LittleEndian(data.Slice(_updatedComponents, 4));
                return data.Slice(_updatedComponents + 4, count * 4);
            }

            public ReadOnlySpan<byte> GetReferenceFrameData(ReadOnlySpan<byte> data)
            {
                // This is a count-prefixed contiguous array of RenderTreeFrame.
                var count = BinaryPrimitives.ReadInt32LittleEndian(data.Slice(_referenceFrames, 4));
                return data.Slice(_referenceFrames + 4, count * ReferenceFrameSize);
            }

            public ReadOnlySpan<byte> GetStringTableIndexes(ReadOnlySpan<byte> data)
            {
                // This is a contiguous array of integers delimited by the end of the data section.
                return data.Slice(_strings, data.Length - 20 - _strings);
            }
        }
    }
}
#nullable restore
