using FrameProcessor.Domain;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.Formats.Png;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using SixLabors.ImageSharp.Processing.Processors.Quantization;

namespace FrameProcessor.ImagePipeline;

public sealed class ImagePipeline : IImagePipeline
{
    public byte[] Process(Stream input, Frame frame)
    {
        ArgumentNullException.ThrowIfNull(input);
        ArgumentNullException.ThrowIfNull(frame);

        var dither = DitheringRegistry.All[frame.Dithering];
        var displayPalette = PaletteFactory.BuildDisplay(frame.Palette);
        var devicePalette = PaletteFactory.BuildDevice(frame.Palette);
        var oriented = frame.Resolution.ForOrientation(frame.Orientation);

        using var image = Image.Load<Rgba32>(input);

        image.Mutate(ctx => ctx
            .Resize(new ResizeOptions
            {
                Size = new Size(oriented.Width, oriented.Height),
                Mode = ResizeMode.Crop,
                Position = AnchorPositionMode.Center,
            })
            .Dither(dither, displayPalette));

        RemapDisplayToDevice(image, frame.Palette);

        if (frame.Orientation == Orientation.Portrait)
        {
            image.Mutate(ctx => ctx.Rotate(RotateMode.Rotate90));
        }

        var encoder = new PngEncoder
        {
            ColorType = PngColorType.Palette,
            BitDepth = SmallestBitDepthFor(frame.Palette.Count),
            Quantizer = new PaletteQuantizer(devicePalette),
        };

        using var output = new MemoryStream();
        image.SaveAsPng(output, encoder);
        return output.ToArray();
    }

    private static void RemapDisplayToDevice(Image<Rgba32> image, IReadOnlyList<PaletteEntry> palette)
    {
        var map = new Dictionary<Rgba32, Rgba32>(palette.Count);
        foreach (var entry in palette)
        {
            map[entry.DisplayColor.ToPixel<Rgba32>()] = entry.DeviceColor.ToPixel<Rgba32>();
        }

        image.ProcessPixelRows(accessor =>
        {
            for (var y = 0; y < accessor.Height; y++)
            {
                var row = accessor.GetRowSpan(y);
                for (var x = 0; x < row.Length; x++)
                {
                    if (map.TryGetValue(row[x], out var mapped))
                    {
                        row[x] = mapped;
                    }
                }
            }
        });
    }

    private static PngBitDepth SmallestBitDepthFor(int paletteCount) => paletteCount switch
    {
        <= 2 => PngBitDepth.Bit1,
        <= 4 => PngBitDepth.Bit2,
        <= 16 => PngBitDepth.Bit4,
        _ => PngBitDepth.Bit8,
    };
}