using System.Runtime.CompilerServices;
using SixLabors.ImageSharp.PixelFormats;
using StitchATon2.Infra;
using StitchATon2.Infra.Buffers;
using StitchATon2.Infra.Encoders;

namespace StitchATon2.Domain.ImageCreators;

public class ImageCreator : IDisposable
{
    private readonly GridSection _section;

    private int FullWidth => _section.TileManager.Configuration.FullWidth;
    private int FullHeight => _section.TileManager.Configuration.FullHeight;

    private int OffsetX => _section.OffsetX;
    private int OffsetY => _section.OffsetY;

    private int Width => _section.Width;
    private int Height => _section.Height;

    private int TileWidth => _section.TileManager.Configuration.Width;
    private int TileHeight => _section.TileManager.Configuration.Height;
    private Tile TileOrigin => _section.Origin;
    
    private int RightmostPixelIndex => _section.TileManager.Configuration.RightTileIndex;
    private int BottomPixelIndex => _section.TileManager.Configuration.BottomTileIndex;
    
    private TileManager TileManager => _section.TileManager;

    private readonly ArrayOwner<Int32Pixel> _mmfReadBuffer;

    public ImageCreator(GridSection section)
    {
        _section = section;
        _mmfReadBuffer = MemoryAllocator.AllocateArray<Int32Pixel>(TileWidth);
    }
    
    public async Task WriteToStream(Stream writableStream, float scale)
    {
        var scaleFactor = MathF.ReciprocalEstimate(scale);
        var targetWidth = (int)(Width / scaleFactor);
        var targetHeight = (int)(Height / scaleFactor);
        
        var encoder = new PngStreamEncoder(writableStream, targetWidth, targetHeight);
        await encoder.WriteHeader();

        var outputBufferSize = targetWidth * Unsafe.SizeOf<Rgb24>();
        using var outputBuffer = MemoryAllocator.AllocateManaged<byte>(outputBufferSize);
        
        using var xLookup = Utils.BoundsMatrix(scaleFactor, targetWidth, FullWidth, OffsetX);
        using var yLookup = Utils.BoundsMatrix(scaleFactor, targetHeight, FullHeight, OffsetY);

        using var yStartMap = MemoryAllocator.Allocate<Int32Pixel>(targetWidth);
        using var yEndMap = MemoryAllocator.Allocate<Int32Pixel>(targetWidth);

        var yStart = OffsetY;
        Task? outputTask = null;
        for (var y = 0; y < targetHeight; y++)
        {
            var yEnd = yLookup[y];
            
            var (localRow0, localOffsetY0) = int.DivRem(yStart, TileHeight);
            MapRow(localRow0, localOffsetY0, xLookup.Span[..targetWidth], yStartMap);
            
            var (localRow1, localOffsetY1) = int.DivRem(yEnd, TileHeight);
            MapRow(localRow1, localOffsetY1, xLookup.Span[..targetWidth], yEndMap);
            
            if (localRow0 != localRow1)
            {
                MapRowAppend(localRow0, BottomPixelIndex, xLookup.Span[..targetWidth], yEndMap);
            }

            if(outputTask != null)
                await outputTask;

            int xStart = OffsetX, x0 = 0;
            for (int x1 = 0, i = 0; x1 < targetWidth; x1++)
            {
                var xEnd = xLookup[x1];

                var pixel = yEndMap[x1];
                pixel += yStartMap[x0];
                pixel -= yEndMap[x0];
                pixel -= yStartMap[x1];
                
                pixel /= Math.Max(1, (xEnd - xStart) * (yEnd - yStart));
                outputBuffer.Memory.Span[i++] = (byte)pixel.R;
                outputBuffer.Memory.Span[i++] = (byte)pixel.G;
                outputBuffer.Memory.Span[i++] = (byte)pixel.B;
                
                xStart = xEnd;
                x0 = x1;
            }
            
            outputTask = encoder.WriteData(outputBuffer.Memory[..outputBufferSize]);
            yStart = yEnd;
        }

        await encoder.WriteEndOfFile();
    }
    
    private void MapRow(int rowOffset, int yOffset, Span<int> sourceMap, IBuffer<Int32Pixel> destination)
    {
        var currentTile = TileManager.GetAdjacent(TileOrigin, 0, rowOffset);
        var xAdder = Int32Pixel.Zero;
        var xOffset = 0;
        var written = 0;
        while (true)
        {
            currentTile.Integral.Acquire(yOffset, _mmfReadBuffer.Array);
            int localX;
            while (written < sourceMap.Length && (localX = sourceMap[written] - xOffset) < TileWidth)
            {
                destination.Span[written] = _mmfReadBuffer[localX];
                destination.Span[written] += xAdder;
                written++;
            }

            if (written >= sourceMap.Length)
                break;

            xAdder += _mmfReadBuffer[RightmostPixelIndex];
            xOffset += TileWidth;
            currentTile = TileManager.GetAdjacent(currentTile, 1, 0);
        }
    }

    private void MapRowAppend(int rowOffset, int yOffset, Span<int> sourceMap, IBuffer<Int32Pixel> destination)
    {
        var currentTile = TileManager.GetAdjacent(TileOrigin, 0, rowOffset);
        var xAdder = Int32Pixel.Zero;
        var xOffset = 0;
        var written = 0;
        while (true)
        {
            currentTile.Integral.Acquire(yOffset, _mmfReadBuffer.Array);
            int localX;
            while (written < sourceMap.Length && (localX = sourceMap[written] - xOffset) < TileWidth)
            {
                destination.Span[written] += _mmfReadBuffer[localX];
                destination.Span[written] += xAdder;
                written++;
            }

            if (written >= sourceMap.Length)
                break;

            xAdder += _mmfReadBuffer[RightmostPixelIndex];
            xOffset += TileWidth;
            currentTile = TileManager.GetAdjacent(currentTile, 1, 0);
        }
    }

    public void Dispose()
    {
        _mmfReadBuffer.Dispose();
        GC.SuppressFinalize(this);
    }
}