using System.Buffers.Binary;

using OpenCvSharp;

using LibCameraSharp;

namespace GSS2.FrameDecoders;

public class BGGR_PISP_COMP1_RAW_FrameDecoder : IFrameDecoder // BGGR_PISP_COMP1/RAW
{
    // Не знаю как оно работает, но работает
    static class PispDecompressor
    {
        const int COMPRESS_OFFSET = 2048;

        public static void DecompressPisp(ReadOnlySpan<byte> src, int width, int height, int stride, Span<ushort> dst)
        {
            int paddedWidth = (width + 7) & ~7;

            for (int y = 0; y < height; y++)
            {
                int srcRow = y * stride;
                int dstRow = y * paddedWidth;

                int sp = srcRow;
                int dp = dstRow;

                for (int x = 0; x < paddedWidth; x += 8)
                {
                    uint w0 = BinaryPrimitives.ReadUInt32LittleEndian(src.Slice(sp, 4));
                    uint w1 = BinaryPrimitives.ReadUInt32LittleEndian(src.Slice(sp + 4, 4));
                    sp += 8;

                    SubBlock(dst, dp, w0); // even pixels
                    SubBlock(dst, dp + 1, w1); // odd pixels

                    for (int i = 0; i < 8; i++)
                    {
                        dst[dp + i] = Postprocess(dst[dp + i]);
                    }

                    dp += 8;
                }
            }
        }

        static ushort Postprocess(ushort a)
        {
            int v = a + COMPRESS_OFFSET;
            return (ushort)Math.Min(0xFFFF, v);
        }

        static void SubBlock(Span<ushort> d, int baseIndex, uint w)
        {
            int qmode = (int)(w & 3);
            int[] q = new int[4];

            if (qmode < 3)
            {
                int field0 = (int)((w >> 2) & 511);
                int field1 = (int)((w >> 11) & 127);
                int field2 = (int)((w >> 18) & 127);
                int field3 = (int)((w >> 25) & 127);

                if (qmode == 2 && field0 >= 384)
                {
                    q[1] = field0;
                    q[2] = field1 + 384;
                }
                else
                {
                    q[1] = (field1 >= 64) ? field0 : field0 + 64 - field1;
                    q[2] = (field1 >= 64) ? field0 + field1 - 64 : field0;
                }

                int p1 = Math.Max(0, q[1] - 64);
                int p2 = Math.Max(0, q[2] - 64);
                if (qmode == 2)
                {
                    p1 = Math.Min(384, p1);
                    p2 = Math.Min(384, p2);
                }

                q[0] = p1 + field2;
                q[3] = p2 + field3;
            }
            else
            {
                int pack0 = (int)((w >> 2) & 32767);
                int pack1 = (int)((w >> 17) & 32767);

                q[0] = (pack0 & 15) + 16 * ((pack0 >> 8) / 11);
                q[1] = (pack0 >> 4) % 176;
                q[2] = (pack1 & 15) + 16 * ((pack1 >> 8) / 11);
                q[3] = (pack1 >> 4) % 176;
            }

            d[baseIndex + 0] = Dequantize(q[0], qmode);
            d[baseIndex + 2] = Dequantize(q[1], qmode);
            d[baseIndex + 4] = Dequantize(q[2], qmode);
            d[baseIndex + 6] = Dequantize(q[3], qmode);
        }

        static ushort Dequantize(int q, int qmode)
        {
            int v = qmode switch
            {
                0 => (q < 320) ? 16 * q : 32 * (q - 160),
                1 => 64 * q,
                2 => 128 * q,
                _ => (q < 94) ? 256 * q : Math.Min(0xFFFF, 512 * (q - 47))
            };

            return (ushort)Math.Min(0xFFFF, v);
        }
    }

    public uint Fourcc => 0x32525942; // BYR2 
    // Модификатор 0xC00000000000001 в соответствии с libcamera является PISP_FORMAT_MOD_COMPRESS_MODE1
    // https://github.com/raspberrypi/libcamera/blob/f0e40f1c50bd0afe65727d6e407d0dcb42666ada/include/linux/drm_fourcc.h#L1692
    // Режимы компрессии Raspberry Pi PiSP описаны в документации ядра linux:
    // 2.6.1.3.1. Raspberry Pi PiSP compressed 8-bit Bayer formats
    // https://docs.kernel.org/userspace-api/media/v4l/pixfmt-srggb8-pisp-comp.html
    // Пример реализации декодера приведён в репозитории https://github.com/raspberrypi/rpicam-apps
    // https://github.com/raspberrypi/rpicam-apps/blob/d821489c2ec89f541915393a823555632653a58f/image/dng.cpp
    public ulong Modifier => 0xC00000000000001; // PISP_FORMAT_MOD_COMPRESS_MODE1
    public ColorSpace.PrimariesEnum PrimariesEnum => ColorSpace.PrimariesEnum.Raw;
    public ColorSpace.TransferFunctionEnum TransferFunctionEnum => ColorSpace.TransferFunctionEnum.Linear;
    public ColorSpace.YcbcrEncodingEnum YcbcrEncodingEnum => ColorSpace.YcbcrEncodingEnum.None;
    public ColorSpace.RangeEnum RangeEnum => ColorSpace.RangeEnum.Full;

    public Mat Decode(int width, int height, int stride, List<byte[]> planes, double rGain, double gGain, double bGain, int blackLevel)
    {
        // Декодирование этого формата состоит из двух этапов:
        // 1) Распаковка исходных данных в которых каждые 8 байт представляют 8 пикселей (это не значит что 1 байт представляет 1 пиксель!)
        // 2) Демозаика полученного Bayer изображения в RGB (BGR)

        if (planes.Count < 1)
            throw new ArgumentException($"At least 1 plane expected, got {planes.Count}");

        var pixelCount = width * height;

        var unpackedSum = new uint[pixelCount];
        foreach (var source in planes)
        {
            var unpacked = new ushort[pixelCount];
            PispDecompressor.DecompressPisp(source, width, height, stride, unpacked);
            for (int i = 0; i < pixelCount; i++)
                unpackedSum[i] += unpacked[i];
        }

        var unpackedAvg = new ushort[pixelCount];
        for (int i = 0; i < pixelCount; i++)
            unpackedAvg[i] = (ushort)(unpackedSum[i] / planes.Count());

        for (int y = 0; y < height; y++)
            for (int x = 0; x < width; x++)
            {
                int idx = y * width + x;
                double value = unpackedAvg[idx];

                value = Math.Max(0, value - blackLevel);

                bool isBlue = (y % 2 == 0) && (x % 2 == 0);
                bool isRed = (y % 2 == 1) && (x % 2 == 1);

                if (isRed)
                    value *= rGain;
                else if (isBlue)
                    value *= bGain;
                else
                    value *= gGain;

                unpackedAvg[idx] = (ushort)Math.Min(value, ushort.MaxValue);
            }

        using var bayer = Mat.FromPixelData(height, width, MatType.CV_16UC1, unpackedAvg);
        var bgr = new Mat();
        Cv2.CvtColor(bayer, bgr, ColorConversionCodes.BayerRG2BGR);

        return bgr;
    }
}