using System.Numerics;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;

/* 
 * Index upscaler code (converted/reversed for downscaling purposes) provided by Ny
 * thank you!!
*/

namespace LightlessSync.Services.TextureCompression;

internal static class IndexDownscaler
{
    private static readonly Vector2[] SampleOffsets =
    {
        new(0.25f, 0.25f),
        new(0.75f, 0.25f),
        new(0.25f, 0.75f),
        new(0.75f, 0.75f),
    };

    public static Image<Rgba32> Downscale(Image<Rgba32> source, int targetWidth, int targetHeight, int blockMultiple)
    {
        var current = source.Clone();

        while (current.Width > targetWidth || current.Height > targetHeight)
        {
            var nextWidth = Math.Max(targetWidth, Math.Max(blockMultiple, current.Width / 2));
            var nextHeight = Math.Max(targetHeight, Math.Max(blockMultiple, current.Height / 2));
            var next = new Image<Rgba32>(nextWidth, nextHeight);

            for (var y = 0; y < nextHeight; y++)
            {
                var srcY = Math.Min(current.Height - 1, y * 2);
                for (var x = 0; x < nextWidth; x++)
                {
                    var srcX = Math.Min(current.Width - 1, x * 2);

                    var topLeft = current[srcX, srcY];
                    var topRight = current[Math.Min(current.Width - 1, srcX + 1), srcY];
                    var bottomLeft = current[srcX, Math.Min(current.Height - 1, srcY + 1)];
                    var bottomRight = current[Math.Min(current.Width - 1, srcX + 1), Math.Min(current.Height - 1, srcY + 1)];

                    next[x, y] = DownscaleIndexBlock(topLeft, topRight, bottomLeft, bottomRight);
                }
            }

            current.Dispose();
            current = next;
        }

        return current;
    }

    private static Rgba32 DownscaleIndexBlock(in Rgba32 topLeft, in Rgba32 topRight, in Rgba32 bottomLeft, in Rgba32 bottomRight)
    {
        Span<Rgba32> ordered = stackalloc Rgba32[4]
        {
            bottomLeft,
            bottomRight,
            topRight,
            topLeft
        };

        Span<float> weights = stackalloc float[4];
        var hasContribution = false;

        foreach (var sample in SampleOffsets)
        {
            if (TryAccumulateSampleWeights(ordered, sample, weights))
            {
                hasContribution = true;
            }
        }

        if (hasContribution)
        {
            var bestIndex = IndexOfMax(weights);
            if (bestIndex >= 0 && weights[bestIndex] > 0f)
            {
                return ordered[bestIndex];
            }
        }

        Span<Rgba32> fallback = stackalloc Rgba32[4] { topLeft, topRight, bottomLeft, bottomRight };
        return PickMajorityColor(fallback);
    }

    private static bool TryAccumulateSampleWeights(ReadOnlySpan<Rgba32> colors, in Vector2 sampleUv, Span<float> weights)
    {
        var red = new Vector4(
            colors[0].R / 255f,
            colors[1].R / 255f,
            colors[2].R / 255f,
            colors[3].R / 255f);

        var symbols = QuantizeSymbols(red);
        var cellUv = ComputeShiftedUv(sampleUv);

        Span<int> order = stackalloc int[4];
        order[0] = 0;
        order[1] = 1;
        order[2] = 2;
        order[3] = 3;

        ApplySymmetry(ref symbols, ref cellUv, order);

        var equality = BuildEquality(symbols, symbols.W);
        var selector = BuildSelector(equality, symbols, cellUv);

        const uint lut = 0x00000C07u;

        if (((lut >> (int)selector) & 1u) != 0u)
        {
            weights[order[3]] += 1f;
            return true;
        }

        if (selector == 3u)
        {
            equality = BuildEquality(symbols, symbols.Z);
        }

        var weight = ComputeWeight(equality, cellUv);
        if (weight <= 1e-6f)
        {
            return false;
        }

        var factor = 1f / weight;

        var wW = equality.W * (1f - cellUv.X) * (1f - cellUv.Y) * factor;
        var wX = equality.X * (1f - cellUv.X) * cellUv.Y * factor;
        var wZ = equality.Z * cellUv.X * (1f - cellUv.Y) * factor;
        var wY = equality.Y * cellUv.X * cellUv.Y * factor;

        var contributed = false;

        if (wW > 0f)
        {
            weights[order[3]] += wW;
            contributed = true;
        }

        if (wX > 0f)
        {
            weights[order[0]] += wX;
            contributed = true;
        }

        if (wZ > 0f)
        {
            weights[order[2]] += wZ;
            contributed = true;
        }

        if (wY > 0f)
        {
            weights[order[1]] += wY;
            contributed = true;
        }

        return contributed;
    }

    private static Vector4 QuantizeSymbols(in Vector4 channel)
        => new(
            Quantize(channel.X),
            Quantize(channel.Y),
            Quantize(channel.Z),
            Quantize(channel.W));

    private static float Quantize(float value)
    {
        var clamped = Math.Clamp(value, 0f, 1f);
        return (MathF.Round(clamped * 16f) + 0.5f) / 16f;
    }

    private static void ApplySymmetry(ref Vector4 symbols, ref Vector2 cellUv, Span<int> order)
    {
        if (cellUv.X >= 0.5f)
        {
            symbols = SwapYxwz(symbols, order);
            cellUv.X = 1f - cellUv.X;
        }

        if (cellUv.Y >= 0.5f)
        {
            symbols = SwapWzyx(symbols, order);
            cellUv.Y = 1f - cellUv.Y;
        }
    }

    private static Vector4 BuildEquality(in Vector4 symbols, float reference)
        => new(
            AreEqual(symbols.X, reference) ? 1f : 0f,
            AreEqual(symbols.Y, reference) ? 1f : 0f,
            AreEqual(symbols.Z, reference) ? 1f : 0f,
            AreEqual(symbols.W, reference) ? 1f : 0f);

    private static uint BuildSelector(in Vector4 equality, in Vector4 symbols, in Vector2 cellUv)
    {
        uint selector = 0;
        if (equality.X > 0.5f) selector |= 4u;
        if (equality.Y > 0.5f) selector |= 8u;
        if (equality.Z > 0.5f) selector |= 16u;
        if (AreEqual(symbols.X, symbols.Z)) selector |= 2u;
        if (cellUv.X + cellUv.Y >= 0.5f) selector |= 1u;

        return selector;
    }

    private static float ComputeWeight(in Vector4 equality, in Vector2 cellUv)
        => equality.W * (1f - cellUv.X) * (1f - cellUv.Y)
         + equality.X * (1f - cellUv.X) * cellUv.Y
         + equality.Z * cellUv.X * (1f - cellUv.Y)
         + equality.Y * cellUv.X * cellUv.Y;

    private static Vector2 ComputeShiftedUv(in Vector2 uv)
    {
        var shifted = new Vector2(
            uv.X - MathF.Floor(uv.X),
            uv.Y - MathF.Floor(uv.Y));

        shifted.X -= 0.5f;
        if (shifted.X < 0f)
        {
            shifted.X += 1f;
        }

        shifted.Y -= 0.5f;
        if (shifted.Y < 0f)
        {
            shifted.Y += 1f;
        }

        return shifted;
    }

    private static Vector4 SwapYxwz(in Vector4 v, Span<int> order)
    {
        var o0 = order[0];
        var o1 = order[1];
        var o2 = order[2];
        var o3 = order[3];

        order[0] = o1;
        order[1] = o0;
        order[2] = o3;
        order[3] = o2;

        return new Vector4(v.Y, v.X, v.W, v.Z);
    }

    private static Vector4 SwapWzyx(in Vector4 v, Span<int> order)
    {
        var o0 = order[0];
        var o1 = order[1];
        var o2 = order[2];
        var o3 = order[3];

        order[0] = o3;
        order[1] = o2;
        order[2] = o1;
        order[3] = o0;

        return new Vector4(v.W, v.Z, v.Y, v.X);
    }

    private static int IndexOfMax(ReadOnlySpan<float> values)
    {
        var bestIndex = -1;
        var bestValue = 0f;

        for (var i = 0; i < values.Length; i++)
        {
            if (values[i] > bestValue)
            {
                bestValue = values[i];
                bestIndex = i;
            }
        }

        return bestIndex;
    }

    private static bool AreEqual(float a, float b) => MathF.Abs(a - b) <= 1e-5f;

    private static Rgba32 PickMajorityColor(ReadOnlySpan<Rgba32> colors)
    {
        var counts = new Dictionary<Rgba32, int>(colors.Length);
        foreach (var color in colors)
        {
            if (counts.TryGetValue(color, out var count))
            {
                counts[color] = count + 1;
            }
            else
            {
                counts[color] = 1;
            }
        }

        return counts
            .OrderByDescending(kvp => kvp.Value)
            .ThenByDescending(kvp => kvp.Key.A)
            .ThenByDescending(kvp => kvp.Key.R)
            .ThenByDescending(kvp => kvp.Key.G)
            .ThenByDescending(kvp => kvp.Key.B)
            .First().Key;
    }
}