using System;
using System.Collections.Generic;
using System.IO;
using Minint.Core.Services;

namespace Minint.Infrastructure.Export;

/// <summary>
/// Self-implemented GIF89a animated exporter.
/// Quantizes each ARGB frame to 256 colors (including transparent) using popularity.
/// Uses LZW compression as required by the GIF spec.
/// </summary>
public sealed class GifExporter : IGifExporter
{
    public void Export(Stream stream, IReadOnlyList<(uint[] Pixels, uint DelayMs)> frames, int width, int height)
    {
        ArgumentNullException.ThrowIfNull(stream);
        if (frames.Count == 0)
            throw new ArgumentException("At least one frame is required.");

        using var w = new BinaryWriter(stream, System.Text.Encoding.UTF8, leaveOpen: true);

        WriteGifHeader(w, width, height);
        WriteNetscapeExtension(w); // infinite loop

        foreach (var (pixels, delayMs) in frames)
        {
            var (palette, indices, transparentIndex) = Quantize(pixels, width * height);
            int colorBits = GetColorBits(palette.Length);
            int tableSize = 1 << colorBits;

            WriteGraphicControlExtension(w, delayMs, transparentIndex);
            WriteImageDescriptor(w, width, height, colorBits);
            WriteColorTable(w, palette, tableSize);
            WriteLzwImageData(w, indices, colorBits);
        }

        w.Write((byte)0x3B); // GIF trailer
    }

    private static void WriteGifHeader(BinaryWriter w, int width, int height)
    {
        w.Write("GIF89a"u8.ToArray());
        w.Write((ushort)width);
        w.Write((ushort)height);
        // Global color table flag=0, color resolution=7, sort=0, gct size=0
        w.Write((byte)0x70);
        w.Write((byte)0); // background color index
        w.Write((byte)0); // pixel aspect ratio
    }

    private static void WriteNetscapeExtension(BinaryWriter w)
    {
        w.Write((byte)0x21); // extension introducer
        w.Write((byte)0xFF); // application extension
        w.Write((byte)11); // block size
        w.Write("NETSCAPE2.0"u8.ToArray());
        w.Write((byte)3); // sub-block size
        w.Write((byte)1); // sub-block ID
        w.Write((ushort)0); // loop count: 0 = infinite
        w.Write((byte)0); // block terminator
    }

    private static void WriteGraphicControlExtension(BinaryWriter w, uint delayMs, int transparentIndex)
    {
        w.Write((byte)0x21); // extension introducer
        w.Write((byte)0xF9); // graphic control label
        w.Write((byte)4); // block size
        byte packed = (byte)(transparentIndex >= 0 ? 0x09 : 0x08);
        // disposal method=2 (restore to background), no user input, transparent flag
        w.Write(packed);
        ushort delayCs = (ushort)(delayMs / 10); // GIF delay is in centiseconds
        if (delayCs == 0 && delayMs > 0) delayCs = 1;
        w.Write(delayCs);
        w.Write((byte)(transparentIndex >= 0 ? transparentIndex : 0));
        w.Write((byte)0); // block terminator
    }

    private static void WriteImageDescriptor(BinaryWriter w, int width, int height, int colorBits)
    {
        w.Write((byte)0x2C); // image separator
        w.Write((ushort)0); // left
        w.Write((ushort)0); // top
        w.Write((ushort)width);
        w.Write((ushort)height);
        byte packed = (byte)(0x80 | (colorBits - 1)); // local color table, not interlaced
        w.Write(packed);
    }

    private static void WriteColorTable(BinaryWriter w, byte[][] palette, int tableSize)
    {
        for (int i = 0; i < tableSize; i++)
        {
            if (i < palette.Length)
            {
                w.Write(palette[i][0]); // R
                w.Write(palette[i][1]); // G
                w.Write(palette[i][2]); // B
            }
            else
            {
                w.Write((byte)0);
                w.Write((byte)0);
                w.Write((byte)0);
            }
        }
    }

    #region LZW compression

    private static void WriteLzwImageData(BinaryWriter w, byte[] indices, int colorBits)
    {
        int minCodeSize = Math.Max(colorBits, 2);
        w.Write((byte)minCodeSize);

        var output = new List<byte>();
        LzwCompress(indices, minCodeSize, output);

        int offset = 0;
        while (offset < output.Count)
        {
            int blockLen = Math.Min(255, output.Count - offset);
            w.Write((byte)blockLen);
            for (int i = 0; i < blockLen; i++)
                w.Write(output[offset + i]);
            offset += blockLen;
        }
        w.Write((byte)0); // block terminator
    }

    private static void LzwCompress(byte[] indices, int minCodeSize, List<byte> output)
    {
        int clearCode = 1 << minCodeSize;
        int eoiCode = clearCode + 1;

        int codeSize = minCodeSize + 1;
        int nextCode = eoiCode + 1;
        int maxCode = (1 << codeSize) - 1;

        var table = new Dictionary<(int Prefix, byte Suffix), int>();

        int bitBuffer = 0;
        int bitCount = 0;

        void EmitCode(int code)
        {
            bitBuffer |= code << bitCount;
            bitCount += codeSize;
            while (bitCount >= 8)
            {
                output.Add((byte)(bitBuffer & 0xFF));
                bitBuffer >>= 8;
                bitCount -= 8;
            }
        }

        void ResetTable()
        {
            table.Clear();
            codeSize = minCodeSize + 1;
            nextCode = eoiCode + 1;
            maxCode = (1 << codeSize) - 1;
        }

        EmitCode(clearCode);
        ResetTable();

        if (indices.Length == 0)
        {
            EmitCode(eoiCode);
            if (bitCount > 0) output.Add((byte)(bitBuffer & 0xFF));
            return;
        }

        int prefix = indices[0];

        for (int i = 1; i < indices.Length; i++)
        {
            byte suffix = indices[i];
            var key = (prefix, suffix);

            if (table.TryGetValue(key, out int existing))
            {
                prefix = existing;
            }
            else
            {
                EmitCode(prefix);

                if (nextCode <= 4095)
                {
                    table[key] = nextCode++;
                    if (nextCode > maxCode + 1 && codeSize < 12)
                    {
                        codeSize++;
                        maxCode = (1 << codeSize) - 1;
                    }
                }
                else
                {
                    EmitCode(clearCode);
                    ResetTable();
                }

                prefix = suffix;
            }
        }

        EmitCode(prefix);
        EmitCode(eoiCode);
        if (bitCount > 0) output.Add((byte)(bitBuffer & 0xFF));
    }

    #endregion

    #region Quantization

    /// <summary>
    /// Quantizes ARGB pixels to max 256 palette entries.
    /// Reserves index 0 for transparent if any pixel has alpha &lt; 128.
    /// Uses popularity-based selection for opaque colors.
    /// </summary>
    private static (byte[][] Palette, byte[] Indices, int TransparentIndex) Quantize(uint[] argb, int count)
    {
        bool hasTransparency = false;
        var colorCounts = new Dictionary<uint, int>();

        for (int i = 0; i < count; i++)
        {
            uint px = argb[i];
            byte a = (byte)(px >> 24);
            if (a < 128)
            {
                hasTransparency = true;
                continue;
            }
            uint opaque = px | 0xFF000000u;
            colorCounts.TryGetValue(opaque, out int c);
            colorCounts[opaque] = c + 1;
        }

        int transparentIndex = hasTransparency ? 0 : -1;
        int maxColors = hasTransparency ? 255 : 256;

        var sorted = new List<KeyValuePair<uint, int>>(colorCounts);
        sorted.Sort((a, b) => b.Value.CompareTo(a.Value));
        if (sorted.Count > maxColors)
            sorted.RemoveRange(maxColors, sorted.Count - maxColors);

        var palette = new List<byte[]>();
        var colorToIndex = new Dictionary<uint, byte>();

        if (hasTransparency)
        {
            palette.Add([0, 0, 0]); // transparent slot
        }

        foreach (var kv in sorted)
        {
            uint px = kv.Key;
            byte idx = (byte)palette.Count;
            palette.Add([
                (byte)((px >> 16) & 0xFF),
                (byte)((px >> 8) & 0xFF),
                (byte)(px & 0xFF)
            ]);
            colorToIndex[px] = idx;
        }

        if (palette.Count == 0)
            palette.Add([0, 0, 0]);

        var indices = new byte[count];
        for (int i = 0; i < count; i++)
        {
            uint px = argb[i];
            byte a = (byte)(px >> 24);
            if (a < 128)
            {
                indices[i] = (byte)(transparentIndex >= 0 ? transparentIndex : 0);
                continue;
            }

            uint opaque = px | 0xFF000000u;
            if (colorToIndex.TryGetValue(opaque, out byte idx))
            {
                indices[i] = idx;
            }
            else
            {
                indices[i] = FindClosest(opaque, palette, hasTransparency ? 1 : 0);
            }
        }

        return (palette.ToArray(), indices, transparentIndex);
    }

    private static byte FindClosest(uint argb, List<byte[]> palette, int startIdx)
    {
        byte r = (byte)((argb >> 16) & 0xFF);
        byte g = (byte)((argb >> 8) & 0xFF);
        byte b = (byte)(argb & 0xFF);

        int bestIdx = startIdx;
        int bestDist = int.MaxValue;
        for (int i = startIdx; i < palette.Count; i++)
        {
            int dr = r - palette[i][0];
            int dg = g - palette[i][1];
            int db = b - palette[i][2];
            int dist = dr * dr + dg * dg + db * db;
            if (dist < bestDist)
            {
                bestDist = dist;
                bestIdx = i;
            }
        }
        return (byte)bestIdx;
    }

    private static int GetColorBits(int paletteCount)
    {
        int bits = 2;
        while ((1 << bits) < paletteCount) bits++;
        return Math.Min(bits, 8);
    }

    #endregion
}