BlenderSharp/BlendFile/Reader.cs

using System.Reflection;
using System.Runtime.CompilerServices;
using System.Text;
using CodeGenerator;
using Kaitai;
using static Kaitai.BlendFile;


namespace BlendFile;

public class Reader {
    string _path;
    private readonly Dictionary<int, Type> dnaTypes = new();
    private readonly Dictionary<string, Type> dnaTypesDb = new();

    private Dictionary<(IntPtr, Type), object> objects = new();

    /// <summary>
    /// The objects that have been converted from the blend file
    /// </summary>
    public Dictionary<(IntPtr, Type), object> Objects => objects;

    public Dictionary<IntPtr, object> InstantiatedObjects = new();

    /// <summary>
    /// A dictionary that contains pointers to pointers
    /// </summary>
    /// <remarks>
    /// Key: Memory address of the pointer
    /// Value: Memory address of the object we are pointing to
    /// </remarks>
    private Dictionary<IntPtr, IntPtr> pointers = new();

    /// <summary>
    /// Gets all the converted objects
    /// </summary>
    public List<object> GetObjects() => objects.Values.ToList();

    /// <summary>
    /// Gets all the converted objects of a specific type
    /// </summary>
    /// <typeparam name="T">The type of objects that will be retrieved</typeparam>
    public List<T> GetObjects<T>() => objects.Values.OfType<T>().ToList();

    private SortedDictionary<long, FileBlock> memBlocks = new();

    /// <summary>
    /// Gets the block at the specified memory address
    /// </summary>
    /// <param name="memAddr">memory address in current system endianness</param>
    /// <returns>A <see cref="Kaitai.BlendFile.FileBlock"/> object</returns>
    /*
    public FileBlock? GetBlock(long memAddr) {
        var blocks = OrderedBlocks.Where(x =>
            memAddr >= x.MemAddr.ToPointer() && memAddr < x.MemAddr.ToPointer() + x.LenBody);

        return blocks.MaxBy(x => x.LenBody);
    }
    */

    public FileBlock? GetBlock(long memAddr)
    { 
        return OrderedBlocks.FirstOrDefault(x => x.MemAddr.ToMemAddr() == memAddr);
        
    }
    
    

    /// <summary>
    /// Creates a new instance of the <see cref="Reader"/> class
    /// </summary>
    /// <param name="path">A <see cref="string"/> containing a path to a blend file that will be read.</param>
    public Reader(string path) : this() {
        _path = path;
    }

    public List<FileBlock> OrderedBlocks;

    /// <summary>
    /// Creates a new instance of the <see cref="Reader"/> class
    /// </summary>
    public Reader() {
        _path = "";
        var types = Assembly.GetExecutingAssembly().DefinedTypes;
        foreach (var type in types) {
            var attrib = type.GetCustomAttribute<DNAClassAttribute>();
            if (attrib == null) continue;
            dnaTypes.Add(attrib.OriginalIndex, type);
            dnaTypesDb.Add(attrib.OriginalName, type);
        }
    }

    public void Read(string path) {
        _path = path;
        Read();
    }

    public void Read() {
        var file = new KaitaiStream(_path);
        var blend = new Kaitai.BlendFile(file);


        OrderedBlocks = blend.Blocks.OrderBy(x => x.MemAddr.ToPointer()).ToList();

        var startAddr = OrderedBlocks.Skip(1).SkipWhile(x => x.Code != "DNA1").First().MemAddr.ToPointer();
        var endAddr = OrderedBlocks.Last().MemAddr.ToPointer();

        Console.WriteLine($"Start offset: 0x{startAddr:X}");
        Console.WriteLine($"End offset: 0x{endAddr:X}");
        Console.WriteLine($"Size in MB: {((endAddr - startAddr) / 1024 / 1024)}");


        bool isLe = blend.Hdr.Endian == Endian.Le;
        //TODO: two blocks somehow have the same mem address... this sounds wrong.
        blend.Blocks.ForEach(block => memBlocks.TryAdd(block.MemAddr.ToMemAddr(isLe), block));

        foreach (var block in blend.Blocks) {
            //We need to read all blocks of data regardless of the type
            //Checks if the block has a known SDNA type, meaning that we know how it is structured 
            if (!dnaTypes.ContainsKey((int)block.SdnaIndex)) continue;

            //Get the type of the block
            dnaTypesDb.TryGetValue(block.SdnaStruct.Type, out Type? t);
            if (t == null) continue; //should never happen

            // How many objects are in the block
            var count = block.Count;
            //offset for the next object in the block
            var blockOffset = 0;
            //for each expected object in the block
            for (var i = 0; i < count; i++) {
                //create an instance of type "t" (the dna type we inferred before)  
                var obj = ActivateInstance(block.SdnaStruct.Type);
                if (obj == null) continue; //should never happen

                //fill the object with the data from the block
                FillObject(block, ref obj, t.GetFields(), blockOffset);

                //move the offset to the next object in the block
                blockOffset += t.GetCustomAttribute<DNAClassAttribute>()!.Size;
            }
        }
        for(int i=0; i<Objects.Count; i++) {
            var obj = Objects.ElementAt(i);
            // for each converted object
            // get the fields of the object
            FieldInfo[] fieldInfo = obj.Value.GetType().GetFields();
            // get all fields that are pointers
            var list = fieldInfo.Where(fldInfo =>
                fldInfo.GetCustomAttributes().OfType<DNAFieldAttribute>().FirstOrDefault()?.IsPointer ?? false).ToList();

            list.AddRange(fieldInfo.Where(fldInfo =>
                fldInfo.GetCustomAttributes().OfType<DNAArrayAttribute>().FirstOrDefault()?.IsPointer ?? false));

            // for each pointer field
            foreach (var f in list)
            {
                int fieldIndex = f.GetCustomAttribute<DNAAttribute>()!.OriginalIndex;
                // get the pointer value
                var addr = GetBlockFieldDataOffset(obj.Key.Item1, fieldIndex, fieldInfo);
                // dereference the pointer
                addr = pointers.GetValueOrDefault(addr);
                if (addr == IntPtr.Zero) continue; // null pointer
                // get the object that the pointer is pointing to if we've already converted it
                var newobj = objects.GetValueOrDefault((addr, f.FieldType));
                if (newobj != null)
                    f.SetValue(obj.Value, newobj);
                else // if we haven't converted the object yet, we need to convert it now
                {
                    newobj = ActivateInstance(f.FieldType);
                    if (newobj != null) {
                        try
                        {
                            FillObject(addr, ref newobj, f.FieldType.GetFields());
                        } catch (Exception e)
                        {
                            Console.WriteLine($"Error while filling object {f.FieldType.Name}: {e}");
                            continue;
                        }

                        f.SetValue(obj.Value, newobj);
                    } // should never happen
                    else throw new NotSupportedException($"Type \"{f.FieldType}\" is unknown");
                }
            }
        }

        foreach (var ptr in pointers) {
            var obj = objects.GetValueOrDefault((ptr.Key, typeof(object)));
            if (obj == null) continue;
            InstantiatedObjects.Add(ptr.Key, obj);
        }
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private long GetBlockFieldDataOffset(long blockAddress, int fieldIndex, FieldInfo[] fieldMetadata) =>
        blockAddress + GetFieldDataOffset(fieldIndex, fieldMetadata);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private IntPtr GetBlockFieldDataOffset(IntPtr blockAddress, int fieldIndex, FieldInfo[] fieldMetadata) =>
        new(blockAddress + GetFieldDataOffset(fieldIndex, fieldMetadata));

    /// <summary>
    /// Gets the offset of the data of a field in a block
    /// </summary>
    /// <param name="fieldIndex">index of the field in the structure</param>
    /// <param name="fieldMetadata"><see cref="FieldInfo"/> array of metadata</param>
    /// <returns></returns>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private long GetFieldDataOffset(int fieldIndex, FieldInfo[] fieldMetadata)
    {
        var fi = fieldMetadata.First(x => x.GetCustomAttribute<DNAAttribute>()!.OriginalIndex == fieldIndex);
        //get the attribute from the field
        var attrib = fi.GetCustomAttribute<DNAAttribute>();
        return attrib switch {
            DNAFieldAttribute fieldAttribute => fieldAttribute.MemoryOffset,
            DNAArrayAttribute arrayAttribute => arrayAttribute.MemoryOffset,
            _ => 0
        };
    }

    /// <summary>
    /// Creates an instance of a given type
    /// </summary>
    /// <param name="type">A <see cref="string"/> containing the name of the type to create</param>
    /// <returns>An object of the type specified in the parameter or null</returns>
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private object? ActivateInstance(string type) =>
        dnaTypesDb.TryGetValue(type, out Type? t) ? Activator.CreateInstance(t) : null;

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private object? ActivateInstance(Type type) => Activator.CreateInstance(type);

    /// <summary>
    /// Filles a given object with the data from a block, starting to read it from the specified offset
    /// </summary>
    /// <param name="block">data block from where to read the values</param>
    /// <param name="obj">object of same struct type as the one that is being decoded</param>
    /// <param name="fieldMetadata">Array of <see cref="FieldInfo"/>s containing <see cref="DNAFieldAttribute"/> attributes</param>
    /// <param name="startOffset">offset in bytes from where structure starts in the block</param>
    private void FillObject(FileBlock block, ref object? obj, FieldInfo[] fieldMetadata,
        IntPtr startOffset = 0) {
        if (block.Code == "ENDB") return; // ENDB is a special block that does not contain any data
        foreach (var field in fieldMetadata) {
            object? value;
            //Get the DNAFieldAttribute of the current field
            var attrib = field.GetCustomAttribute<DNAAttribute>();
            switch (attrib) {
                case DNAFieldAttribute fieldAttribute:
                    value = ConvertNormalField(block, fieldAttribute, startOffset);
                    break;
                case DNAArrayAttribute arrayAttribute:
                    value = ConvertArrayField(block, field, arrayAttribute, startOffset);
                    break;
                case DNAListAttribute listAttribute:
                    value = ConvertListField(block, field, listAttribute, startOffset);
                    break;
                default:
                    continue; //should never happen. 
            }

            field.SetValue(obj, value);
        }
        //Add the freshly handled object to the database
        objects.Add((block.MemAddr.ToPointer() + startOffset, obj!.GetType()), obj!);
    }

    /// <summary>
    /// Fills a given object from a block, using a pointer to the start of it's data
    /// </summary>
    /// <param name="ptr">pointer to the start of object</param>
    /// <param name="obj">reference to the object to be filled</param>
    /// <param name="fieldMetadata">field metadata info for that object</param>
    /// <exception cref="Exception">Thrown when the block the pointer is pointing to is not found.</exception>
    private void FillObject(IntPtr ptr, ref object? obj, FieldInfo[] fieldMetadata) {
        var block = GetBlock(ptr.ToInt64());
        if (block == null) throw new($"Block for pointer {ptr.ToInt64():X} not found");

        var blockOffset = ptr.ToInt64() - block.MemAddr.ToPointer();
        FillObject(block, ref obj, fieldMetadata, new IntPtr(blockOffset));
    }

    private object? ConvertNormalField(FileBlock block, DNAFieldAttribute attrib, IntPtr startOffset) {
        //Calculate the offset from where the data of the field starts.
        //Because the order of the fields is not guaranteed we need to compute it each time
        IntPtr offset = attrib.MemoryOffset + startOffset;

        //Grab data size, create a container and copy the data from the block to the container
        int size = attrib.Size;
        var data = new byte[size];
        Array.Copy((byte[])block.Body, offset, data, 0, size);

        //Convert the data to the correct type if it's a base type
        object? value = ConvertFieldData(data, attrib.OriginalType);
        if (value != null) return value;

        //Check if the field is a pointer to another DNA structure
        if (dnaTypesDb.ContainsKey(attrib.OriginalType)) {
            if (!attrib.IsPointer) {
                //It's a structure

                //Create a new instance of the DNA structure type
                object? newObj = ActivateInstance(attrib.OriginalType);
                //should never happen... type is missing?
                if (newObj == null) throw new NotSupportedException($"Type \"{attrib.OriginalType}\" is unknown");
                //Get the information of the fields of the new object
                var fieldInfo = newObj.GetType().GetFields();
                //relative offset to this block for the structure
                IntPtr relAddr = block.MemAddr.ToPointer() + offset;
                //If the object is already created, we can just assign it
                if (objects.TryGetValue((relAddr, newObj.GetType()), out object? o)) return o;

                //Fill the object with the data from the block (this is recursive)
                FillObject(block, ref newObj, fieldInfo, offset);
            } else {
                //It's a pointer
                IntPtr memAddr = data.ToPointer();
                if (memAddr == 0) return null; //nullPointer, no need to store the reference
                pointers.TryAdd(block.MemAddr.ToPointer() + offset, data.ToPointer());
            }
        } else {
            //TODO: handle void types. This gets spammed
            //throw new NotSupportedException($"Unknown type \"{attrib.OriginalType}\"");
        }
        return null;
    }

    private static int[] CalculateArrayIndices(List<int> sizes, int index) {
        int[] indexArray = new int[sizes.Count];
        for (int k = 0; k < sizes.Count; k++) {
            indexArray[k] = index / sizes[(k + 1) ..].Aggregate(1, (a, b) => a * b) % sizes[k];
        }

        return indexArray;
    }

    private Array FillPointersArray(DNAArrayAttribute arrayAttribute, byte[] data,
        Type type, List<int> sizes) {
        //Create the array
        var array = Array.CreateInstance(type, sizes.ToArray());

        for (int i = 0; i < arrayAttribute.Size; i += 8) {
            var itemData = new byte[8];
            Array.Copy(data, i, itemData, 0, 8);
            if (itemData.ToPointer() == IntPtr.Zero) continue;
            object? cellData = ActivateInstance(type);
            FillObject(itemData.ToPointer(), ref cellData, type.GetFields());

            var indexArray = CalculateArrayIndices(sizes, i);

            array.SetValue(cellData, indexArray);
        }

        return array;
    }

    private Array FillNormalTypeArray(DNAArrayAttribute arrayAttribute, byte[] data, FileBlock block, Type type,
        List<int> sizes, int itemLenght, int offset) {
        var array = Array.CreateInstance(type, sizes.ToArray());

        for (int i = 0; i < arrayAttribute.Size; i += itemLenght) {
            var itemData = new byte[itemLenght];
            Array.Copy(data, i, itemData, 0, itemLenght);
            object? cellData = ConvertFieldData(itemData, arrayAttribute.OriginalType);

            var indexArray = CalculateArrayIndices(sizes, i);

            array.SetValue(cellData, indexArray);
        }

        return array;
    }

    private Array FillCustomTypeArray(DNAArrayAttribute arrayAttribute, byte[] data, FileBlock block, Type type,
        List<int> sizes, int itemLenght, int offset) {
        var array = Array.CreateInstance(type, sizes.ToArray());

        for (int i = 0; i < arrayAttribute.Size; i += itemLenght) {
            var itemData = new byte[itemLenght];
            Array.Copy(data, i, itemData, 0, itemLenght);
            object? cellData = ActivateInstance(type);
            FillObject(block, ref cellData, type.GetFields(), offset + i);

            var indexArray = CalculateArrayIndices(sizes, i);

            array.SetValue(cellData, indexArray);
        }

        return array;
    }

    private object? ConvertArrayField(FileBlock block, FieldInfo field, DNAArrayAttribute arrayAttribute, IntPtr startOffset) {
        //Calculate the offset from where the data of the field starts.
        //Because the order of the fields is not guaranteed we need to compute it each time
        IntPtr offset = arrayAttribute.MemoryOffset + startOffset;

        //Grab data size, create a container and copy the data from the block to the container
        var sizes = arrayAttribute.OriginalName.GetArrayDimensions();
        int dataLength = arrayAttribute.Size;
        var data = new byte[dataLength];
        Array.Copy((byte[])block.Body, offset, data, 0, dataLength);

        //Gather Array type
        var type = Type.GetType(arrayAttribute.OriginalType.ParseFType()) ?? dnaTypesDb[arrayAttribute.OriginalType];
        if (type == null) throw new NotSupportedException($"Unknown type \"{arrayAttribute.OriginalType}\"");

        if (arrayAttribute.IsPointer) {
            return FillPointersArray(arrayAttribute, data, type, sizes);
        }

        if (!dnaTypesDb.ContainsKey(arrayAttribute.OriginalType)) {
            return FillNormalTypeArray(arrayAttribute, data, block, type, sizes,
                arrayAttribute.OriginalType.ParseFSize()!.Value, offset.ToInt32());
        }

        int itemLenght = type.GetCustomAttribute<DNAFieldAttribute>()?.Size ??
                         type.GetCustomAttribute<DNAClassAttribute>()!.Size;
        return FillCustomTypeArray(arrayAttribute, data, block, type, sizes, itemLenght, offset.ToInt32());
    }

    private object? ConvertListField(FileBlock block, FieldInfo field, DNAListAttribute attrib, IntPtr startOffset) {
        IntPtr countOffset = attrib.CountMemoryOffset + startOffset;
        IntPtr ptrOffset = attrib.PtrMemoryOffset + startOffset;

        var countLen = attrib.CountFieldName.ParseFSize()!.Value;
        var countData = new byte[countLen];
        Array.Copy((byte[])block.Body, countOffset, countData, 0, countLen);
        var tmpCount = ConvertFieldData(countData, attrib.CountFieldName);
        int count;

        switch (tmpCount) {
            case int i:
                count = i;
                break;
            case short s:
                count = s;
                break;
            default:
                throw new NotSupportedException($"Unknown type \"{attrib.CountFieldName}\"");
        }

        List<Object?> objList = new();

        for (IntPtr ptr = ptrOffset; ptr < ptrOffset + count * sizeof(Int64); ptr += sizeof(Int64)) {
            var data = new byte[sizeof(Int64)];
            Array.Copy((byte[])block.Body, ptr, data, 0, sizeof(Int64));
            IntPtr memAddr = new IntPtr(BitConverter.ToInt64(data, 0));

            if (memAddr == IntPtr.Zero) {
                objList.Add(null);
                continue;
            }

            var obj = ActivateInstance(attrib.UnderlyingType);
            if (obj == null) throw new NotSupportedException($"Type \"{attrib.UnderlyingType}\" is unknown");

            var fieldMetadata = obj.GetType().GetFields();

            FillObject(memAddr, ref obj, fieldMetadata);
            objList.Add(obj);
        }

        return objList;
    }

    private object? ConvertFieldData(byte[] data, string type) {
        return type switch {
            "char" => (char)data[0],
            "short" => BitConverter.ToInt16(data, 0),
            "int" => BitConverter.ToInt32(data, 0),
            "float" => BitConverter.ToSingle(data, 0),
            "double" => BitConverter.ToDouble(data, 0),
            "string" => Encoding.UTF8.GetString(data), // utf8?
            "void" => null, // object?
            "ushort" => BitConverter.ToUInt16(data, 0),
            "uchar" => data[0],
            "int64_t" => BitConverter.ToInt64(data, 0),
            "int8_t" => (sbyte)data[0],
            "uint64_t" => BitConverter.ToUInt64(data, 0),
            _ => null
        };
    }
}