using Ghost.Core;
using Misaki.HighPerformance.LowLevel.Buffer;
using Misaki.HighPerformance.LowLevel.Collections;
using Misaki.HighPerformance.LowLevel.Utilities;
using System.Runtime.CompilerServices;

namespace Ghost.Entities;

internal unsafe struct Chunk : IDisposable
{
    public const int CHUNK_SIZE = 16384; // 16 KB
    public const int BIT_ALIGNMENT = 8;
    public const int BIT_SHIFT = 3; // log2(BIT_ALIGNMENT)
    public const int BIT_ALIGNMENT_MINUS_ONE = BIT_ALIGNMENT - 1;

    private UnsafeArray<byte> _data;
    private int _count;
    private readonly int _capacity;

    public int Count
    {
        readonly get => _count;
        set => _count = value;
    }

    public readonly int Capacity => _capacity;

    public Chunk(int size, int capacity)
    {
        _data = new UnsafeArray<byte>(size, Allocator.Persistent, AllocationOption.Clear);
        _capacity = capacity;
        _count = 0;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly byte* GetUnsafePtr()
    {
        return (byte*)_data.GetUnsafePtr();
    }

    public void Dispose()
    {
        _data.Dispose();
    }
}

internal unsafe struct Archetype : IIdentifierType, IDisposable
{
    internal struct ComponentMemoryLayout
    {
        public int componentID;
        public int size;
        public int offset;
        public int enableBitsOffset; // TODO: Support enableable component
    }

    private struct Edge
    {
        public int componentID;
        public int targetArchetype; // can't use Identifier<Archetype> because cycle causer
    }

    private readonly Identifier<Archetype> _id;
    private readonly Identifier<World> _worldID;

    internal UnsafeBitSet _signature;
    internal UnsafeList<Chunk> _chunks;
    internal UnsafeArray<ComponentMemoryLayout> _layouts;

    private UnsafeArray<int> _componentIDToLayoutIndex;
    // TODO: Is hash map better?
    private UnsafeList<Edge> _edgesAdd;
    private UnsafeList<Edge> _edgesRemove;

    private readonly int _hash;
    private int _entityCapacity;
    private int _maxComponentID;
    private int _entityIdsOffset;

    public readonly Identifier<Archetype> ID => _id;

    public readonly int EntityCapacity => _entityCapacity;
    public readonly int ChunkCount => _chunks.Count;
    public readonly int EntityIDsOffset => _entityIdsOffset;

    public Archetype(Identifier<Archetype> id, Identifier<World> worldID, ReadOnlySpan<Identifier<IComponent>> componentIds)
    {
        _id = id;
        _worldID = worldID;

        _chunks = new UnsafeList<Chunk>(4, Allocator.Persistent);
        _edgesAdd = new UnsafeList<Edge>(4, Allocator.Persistent);
        _edgesRemove = new UnsafeList<Edge>(4, Allocator.Persistent);

        if (componentIds.IsEmpty)
        {
            _signature = new UnsafeBitSet(1, Allocator.Persistent, AllocationOption.Clear);
            _hash = 0;

            _signature.ClearAll();
            _entityCapacity = Chunk.CHUNK_SIZE / sizeof(Entity);

            return;
        }

        var highestComponentID = 0;
        for (var i = 0; i < componentIds.Length; i++)
        {
            if (componentIds[i] > highestComponentID)
            {
                highestComponentID = componentIds[i];
            }
        }

        _signature = new UnsafeBitSet(highestComponentID + 1, Allocator.Persistent, AllocationOption.Clear);
        _hash = _signature.GetHashCode();

        CalculateLayout(componentIds);
    }

    private void CalculateLayout(ReadOnlySpan<Identifier<IComponent>> componentIds)
    {
        var entitySize = sizeof(Entity);
        var entityAlign = (int)MemoryUtility.AlignOf<Entity>();

        var components = (Span<ComponentInfo>)stackalloc ComponentInfo[componentIds.Length];
        for (var i = 0; i < componentIds.Length; i++)
        {
            _signature.SetBit(componentIds[i]);
            components[i] = ComponentRegister.GetComponentInfo(componentIds[i]);
        }

        // Calculate total size per entity to get an initial capacity estimate
        var bytesPerEntity = entitySize;
        var maxComponentID = 0;
        for (var i = 0; i < components.Length; i++)
        {
            var comp = components[i];
            bytesPerEntity += comp.size;
            if (comp.id > maxComponentID)
            {
                maxComponentID = comp.id;
            }
        }

        _maxComponentID = maxComponentID;
        _entityCapacity = Chunk.CHUNK_SIZE / bytesPerEntity;
        _layouts = new UnsafeArray<ComponentMemoryLayout>(components.Length, Allocator.Persistent);
        _componentIDToLayoutIndex = new UnsafeArray<int>(_maxComponentID + 1, Allocator.Persistent);

        _componentIDToLayoutIndex.AsSpan().Fill(-1);

        components.Sort((a, b) => b.alignment.CompareTo(a.alignment));
        var tempOffsets = stackalloc int[components.Length];
        var tempBitmaskOffsets = stackalloc int[components.Length];

        while (_entityCapacity > 0)
        {
            var currentOffset = 0;
            var fits = true;

            currentOffset = (currentOffset + entityAlign - 1) & ~(entityAlign - 1);

            _entityIdsOffset = currentOffset;
            currentOffset += _entityCapacity * entitySize;

            for (var i = 0; i < components.Length; i++)
            {
                var size = components[i].size;
                var align = components[i].alignment;

                currentOffset = (currentOffset + align - 1) & ~(align - 1);
                tempOffsets[i] = currentOffset;
                currentOffset += _entityCapacity * size;

                var bitmaskOffset = -1;
                if (components[i].isEnableable)
                {
                    var bitmaskSize = (_entityCapacity + Chunk.BIT_ALIGNMENT_MINUS_ONE) / Chunk.BIT_ALIGNMENT;
                    // Reserve space for the bitmask (1 bit per entity)

                    currentOffset = (currentOffset + Chunk.BIT_ALIGNMENT_MINUS_ONE) & ~Chunk.BIT_ALIGNMENT_MINUS_ONE; // Align
                    bitmaskOffset = currentOffset;
                    currentOffset += bitmaskSize;
                }

                tempBitmaskOffsets[i] = bitmaskOffset;

                if (currentOffset > Chunk.CHUNK_SIZE)
                {
                    fits = false;
                    break;
                }
            }

            if (fits)
            {
                for (var i = 0; i < components.Length; i++)
                {
                    _layouts[i] = new ComponentMemoryLayout
                    {
                        offset = tempOffsets[i],
                        size = components[i].size,
                        componentID = components[i].id,
                        enableBitsOffset = tempBitmaskOffsets[i],
                    };

                    _componentIDToLayoutIndex[components[i].id] = i;
                }

                return;
            }

            _entityCapacity--;
        }
    }

    public void AllocateEntity(out int chunkIndex, out int rowIndex)
    {
        for (var i = 0; i < _chunks.Count; i++)
        {
            ref var chunk = ref _chunks[i];
            if (chunk.Count < _entityCapacity)
            {
                rowIndex = chunk.Count;
                chunk.Count++;
                chunkIndex = i;

                return;
            }
        }

        // Need to allocate a new chunk
        var newChunk = new Chunk(Chunk.CHUNK_SIZE, _entityCapacity);

        // Set all enable to true by default for enableable components
        for (var i = 0; i < _layouts.Count; i++)
        {
            var layout = _layouts[i];
            if (layout.enableBitsOffset != -1)
            {
                var pChunk = newChunk.GetUnsafePtr();
                var pBits = pChunk + layout.enableBitsOffset;
                MemoryUtility.MemSet(pBits, 0xFF, (nuint)((_entityCapacity + Chunk.BIT_ALIGNMENT_MINUS_ONE) / Chunk.BIT_ALIGNMENT));
            }
        }

        rowIndex = 0;
        newChunk.Count++;
        chunkIndex = _chunks.Count;

        _chunks.Add(newChunk);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly void SetEntity(int chunkIndex, int rowIndex, Entity entity)
    {
        var chunk = _chunks[chunkIndex];
        var chunkBase = chunk.GetUnsafePtr();
        var dst = chunkBase + _entityIdsOffset + (sizeof(Entity) * rowIndex);

        MemoryUtility.MemCpy(&entity, dst, (nuint)sizeof(Entity));
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly ResultStatus SetComponentData(int chunkIndex, int rowIndex, Identifier<IComponent> componentID, void* pComponent)
    {
        var r = GetLayout(componentID);
        if (r.Status != ResultStatus.Success)
        {
            return r.Status;
        }

        var offset = r.Value.offset;
        var chunk = _chunks[chunkIndex];

        var chunkBase = chunk.GetUnsafePtr();
        var size = ComponentRegister.GetComponentInfo(componentID).size;
        var dst = chunkBase + offset + (size * rowIndex);

        MemoryUtility.MemCpy(pComponent, dst, (nuint)size);

        return ResultStatus.Success;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly ResultStatus GetComponentDataPtr(int chunkIndex, int rowIndex, Identifier<IComponent> componentID, void** ppv)
    {
        var r = GetLayout(componentID);
        if (r.Status != ResultStatus.Success)
        {
            return r.Status;
        }

        var offset = r.Value.offset;
        var chunk = _chunks[chunkIndex];

        var chunkBase = chunk.GetUnsafePtr();
        var size = ComponentRegister.GetComponentInfo(componentID).size;
        *ppv = chunkBase + offset + (size * rowIndex);

        return ResultStatus.Success;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public ref Chunk GetChunkReference(int index)
    {
        return ref _chunks[index];
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Result<ComponentMemoryLayout, ResultStatus> GetLayout(int componentID)
    {
        if (componentID >= _componentIDToLayoutIndex.Count)
        {
            return Result.Create(default(ComponentMemoryLayout), ResultStatus.InvalidArgument);
        }

        var layoutIndex = _componentIDToLayoutIndex[componentID];
        if (layoutIndex == -1)
        {
            return Result.Create(default(ComponentMemoryLayout), ResultStatus.NotFound);
        }

        return Result.Create(_layouts[layoutIndex], ResultStatus.Success);
    }

    public ResultStatus RemoveEntity(int chunkIndex, int rowIndex)
    {
        if (chunkIndex < 0 || chunkIndex >= _chunks.Count)
        {
            return ResultStatus.InvalidArgument;
        }

        ref var chunk = ref _chunks[chunkIndex];
        var lastIndex = chunk.Count - 1;

        // If we are NOT removing the very last entity, we must swap.
        if (rowIndex != lastIndex)
        {
            var chunkBase = chunk.GetUnsafePtr();
            var pLastEntity = chunkBase + _entityIdsOffset + (sizeof(Entity) * lastIndex);
            var pRowEntity = chunkBase + _entityIdsOffset + (sizeof(Entity) * rowIndex);

            var wroldResult = World.GetWorld(_worldID);
            if (wroldResult.Status != ResultStatus.Success)
            {
                return wroldResult.Status;
            }

            var result = wroldResult.Value.EntityManager.UpdateEntityLocation(*(Entity*)pLastEntity, _id, chunkIndex, rowIndex);
            if (result != ResultStatus.Success)
            {
                return result;
            }

            // Only operate the swap back after the update is succeed.
            MemoryUtility.MemCpy(pLastEntity, pRowEntity, (nuint)sizeof(Entity));

            for (var i = 0; i <= _layouts.Count; i++)
            {
                var layout = _layouts[i];

                var pRow = chunk.GetUnsafePtr() + layout.offset + (layout.size * rowIndex);
                var pLast = chunk.GetUnsafePtr() + layout.offset + (layout.size * lastIndex);

                MemoryUtility.MemCpy(pLast, pRow, (nuint)layout.size);
            }
        }

        chunk.Count--;
        return ResultStatus.Success;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly bool HasComponent(Identifier<IComponent> componentID)
    {
        return _signature.IsSet(componentID);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public void AddEdgeAdd(Identifier<IComponent> componentID, Identifier<Archetype> targetArchetype)
    {
        _edgesAdd.Add(new Edge
        {
            componentID = componentID,
            targetArchetype = targetArchetype
        });
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Identifier<Archetype> GetEdgeAdd(Identifier<IComponent> componentID)
    {
        for (var i = 0; i < _edgesAdd.Count; i++)
        {
            var edge = _edgesAdd[i];
            if (edge.componentID == componentID)
            {
                return edge.targetArchetype;
            }
        }

        return Identifier<Archetype>.Invalid;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public void AddEdgeRemove(Identifier<IComponent> componentID, Identifier<Archetype> targetArchetype)
    {
        _edgesRemove.Add(new Edge
        {
            componentID = componentID,
            targetArchetype = targetArchetype
        });
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Identifier<Archetype> GetEdgeRemove(Identifier<IComponent> componentID)
    {
        for (var i = 0; i < _edgesRemove.Count; i++)
        {
            var edge = _edgesRemove[i];
            if (edge.componentID == componentID)
            {
                return edge.targetArchetype;
            }
        }

        return Identifier<Archetype>.Invalid;
    }

    public override readonly int GetHashCode()
    {
        return _hash;
    }

    public void Dispose()
    {
        if (_chunks.IsCreated)
        {
            foreach (ref var chunk in _chunks)
            {
                chunk.Dispose();
            }
        }

        _signature.Dispose();
        _chunks.Dispose();
        _componentIDToLayoutIndex.Dispose();
        _layouts.Dispose();
        _edgesAdd.Dispose();
        _edgesRemove.Dispose();
    }
}