using Microsoft.Xna.Framework;

namespace MrGameEng.Graphics;

/// <summary>
/// Writer for one ECS chunk during parallel submission: appends accepted sprites compactly
/// into the chunk's scratch range and counts culled ones. Used by exactly one thread.
/// </summary>
public struct SpriteChunkWriter
{
    private readonly SpriteInstance[] _instances;
    private readonly ulong[] _keys;
    private readonly int _offset;

    /// <summary>Sprites accepted into this chunk's range.</summary>
    public int Count { get; private set; }

    /// <summary>Sprites rejected by culling in this chunk.</summary>
    public int Culled { get; private set; }

    internal SpriteChunkWriter(SpriteInstance[] instances, ulong[] keys, int offset)
    {
        _instances = instances;
        _keys = keys;
        _offset = offset;
        Count = 0;
        Culled = 0;
    }

    /// <summary>Appends one accepted sprite.</summary>
    public void Add(in SpriteInstance instance, ulong sortKey)
    {
        var index = _offset + Count;
        _instances[index] = instance;
        _keys[index] = sortKey;
        Count++;
    }

    /// <summary>Counts one culled sprite.</summary>
    public void AddCulled() => Culled++;
}

/// <summary>One sprite queued for rendering this frame.</summary>
public struct SpriteInstance
{
    /// <summary>Texture region to draw. Never null for submitted instances.</summary>
    public Texture2DRegion Region;

    /// <summary>World-space (or screen-space) center of the quad.</summary>
    public Vector2 Center;

    /// <summary>Half extents after scaling, in pixels. May be negative for negative scale.</summary>
    public Vector2 HalfSize;

    /// <summary>Rotation in radians, clockwise.</summary>
    public float Rotation;

    /// <summary>Tint color.</summary>
    public Color Color;

    /// <summary>Mirroring flags.</summary>
    public SpriteFlip Flip;

    /// <summary>Render layer the instance belongs to.</summary>
    public byte Layer;
}

/// <summary>
/// CPU side of the renderer: collects <see cref="SpriteInstance"/>s with their sort keys
/// and orders them layer → depth → texture using a stable LSD radix sort — sprites with
/// equal keys keep their submission order across frames (no flicker), and sorting stays
/// O(n) on large counts. Allocation-free after warm-up (arrays grow geometrically and are
/// reused across frames).
/// </summary>
public sealed class SpriteBatcher
{
    private const int RadixBits = 16;
    private const int RadixSize = 1 << RadixBits;

    private SpriteInstance[] _instances;
    private ulong[] _keys;
    private ulong[] _keysTemp;
    private int[] _order;
    private int[] _orderTemp;
    private readonly int[] _histogram = new int[RadixSize];
    private int _count;

    /// <summary>Creates a batcher with the given initial capacity.</summary>
    public SpriteBatcher(int initialCapacity = 2048)
    {
        _instances = new SpriteInstance[initialCapacity];
        _keys = new ulong[initialCapacity];
        _keysTemp = new ulong[initialCapacity];
        _order = new int[initialCapacity];
        _orderTemp = new int[initialCapacity];
    }

    /// <summary>Number of sprites submitted this frame.</summary>
    public int Count => _count;

    /// <summary>Queues one sprite.</summary>
    public void Submit(in SpriteInstance instance, ulong sortKey)
    {
        if (_count == _instances.Length)
        {
            Grow();
        }

        _instances[_count] = instance;
        _keys[_count] = sortKey;
        _count++;
    }

    /// <summary>
    /// Sorts all submitted sprites (stable: equal keys keep submission order) and returns
    /// their indices in draw order. Valid until the next <see cref="Clear"/>.
    /// </summary>
    public ReadOnlySpan<int> Sort()
    {
        var n = _count;
        for (var i = 0; i < n; i++)
        {
            _order[i] = i;
        }

        if (n < 2)
        {
            return _order.AsSpan(0, n);
        }

        // Биты, различающиеся хотя бы у одной пары ключей: проходы по одинаковым
        // разрядам (один слой, одна глубина) пропускаются целиком.
        ulong orBits = 0,
            andBits = ~0UL;
        for (var i = 0; i < n; i++)
        {
            orBits |= _keys[i];
            andBits &= _keys[i];
        }

        var differing = orBits ^ andBits;
        var keys = _keys;
        var order = _order;
        var keysOut = _keysTemp;
        var orderOut = _orderTemp;

        for (var shift = 0; shift < 64; shift += RadixBits)
        {
            if ((differing >> shift & (RadixSize - 1)) == 0)
            {
                continue;
            }

            Array.Clear(_histogram, 0, RadixSize);
            for (var i = 0; i < n; i++)
            {
                _histogram[(int)(keys[i] >> shift & (RadixSize - 1))]++;
            }

            var running = 0;
            for (var digit = 0; digit < RadixSize; digit++)
            {
                var bucket = _histogram[digit];
                _histogram[digit] = running;
                running += bucket;
            }

            for (var i = 0; i < n; i++)
            {
                var position = _histogram[(int)(keys[i] >> shift & (RadixSize - 1))]++;
                keysOut[position] = keys[i];
                orderOut[position] = order[i];
            }

            (keys, keysOut) = (keysOut, keys);
            (order, orderOut) = (orderOut, order);
        }

        _keys = keys;
        _keysTemp = keysOut;
        _order = order;
        _orderTemp = orderOut;
        return _order.AsSpan(0, n);
    }

    /// <summary>Returns the instance at <paramref name="index"/> (an index from <see cref="Sort"/>).</summary>
    public ref readonly SpriteInstance this[int index] => ref _instances[index];

    /// <summary>The sorted index array after <see cref="Sort"/> (first <see cref="Count"/> entries are valid).</summary>
    internal int[] SortedOrder => _order;

    /// <summary>Resets the batcher for the next frame. Keeps allocated capacity.</summary>
    public void Clear() => _count = 0;

    // --- Параллельная по-чанковая подача -------------------------------------------------
    // Каждый чанк ECS пишет в свой непересекающийся диапазон scratch-арены, затем диапазоны
    // сливаются в порядке чанков — порядок детерминирован, стабильность сортировки сохраняется.

    private SpriteInstance[] _scratchInstances = [];
    private ulong[] _scratchKeys = [];
    private int[] _chunkOffsets = [];
    private int[] _chunkVisible = [];
    private int[] _chunkCulled = [];
    private int _chunkCount;

    /// <summary>Total culled count reported by chunk writers in the last <see cref="CommitChunks"/>.</summary>
    public int LastChunkCulled { get; private set; }

    /// <summary>
    /// Prepares the scratch arena for chunked submission. <paramref name="chunkLengths"/> are
    /// the entity counts of each ECS chunk, in iteration order.
    /// </summary>
    public void BeginChunks(ReadOnlySpan<int> chunkLengths)
    {
        _chunkCount = chunkLengths.Length;
        if (_chunkOffsets.Length < _chunkCount)
        {
            Array.Resize(ref _chunkOffsets, _chunkCount);
            Array.Resize(ref _chunkVisible, _chunkCount);
            Array.Resize(ref _chunkCulled, _chunkCount);
        }

        var total = 0;
        for (var i = 0; i < _chunkCount; i++)
        {
            _chunkOffsets[i] = total;
            total += chunkLengths[i];
        }

        if (_scratchInstances.Length < total)
        {
            _scratchInstances = new SpriteInstance[total];
            _scratchKeys = new ulong[total];
        }
    }

    /// <summary>Returns the writer for chunk <paramref name="chunkIndex"/>. Each chunk is written by one thread.</summary>
    public SpriteChunkWriter GetChunkWriter(int chunkIndex) =>
        new(_scratchInstances, _scratchKeys, _chunkOffsets[chunkIndex]);

    /// <summary>Records the writer's results. Called by the same thread that filled the writer.</summary>
    public void EndChunk(int chunkIndex, in SpriteChunkWriter writer)
    {
        _chunkVisible[chunkIndex] = writer.Count;
        _chunkCulled[chunkIndex] = writer.Culled;
    }

    /// <summary>
    /// Merges all chunk ranges into the main arrays in chunk order (deterministic, keeps
    /// sort stability) and returns the number of accepted sprites.
    /// </summary>
    public int CommitChunks()
    {
        var visible = 0;
        var culled = 0;
        for (var i = 0; i < _chunkCount; i++)
        {
            visible += _chunkVisible[i];
            culled += _chunkCulled[i];
        }

        while (_count + visible > _instances.Length)
        {
            Grow();
        }

        for (var i = 0; i < _chunkCount; i++)
        {
            var length = _chunkVisible[i];
            if (length == 0)
            {
                continue;
            }

            Array.Copy(_scratchInstances, _chunkOffsets[i], _instances, _count, length);
            Array.Copy(_scratchKeys, _chunkOffsets[i], _keys, _count, length);
            _count += length;
        }

        LastChunkCulled = culled;
        _chunkCount = 0;
        return visible;
    }

    private void Grow()
    {
        var capacity = _instances.Length * 2;
        Array.Resize(ref _instances, capacity);
        Array.Resize(ref _keys, capacity);
        Array.Resize(ref _keysTemp, capacity);
        Array.Resize(ref _order, capacity);
        Array.Resize(ref _orderTemp, capacity);
    }
}