package org.kne.codec.kif;

import java.awt.image.BufferedImage;
import java.awt.image.DataBufferInt;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

import org.kne.membandboost.MembandBoost;

/**
 * int[] 像素数组与 byte[] 的相互转换工具
 * 每个像素按 R, G, B 三个字节依次写入
 */
public class PixelSerializer {

    // ==================== 基础序列化/反序列化 (RGB, 3通道) ====================

    public static byte[] serialize(int[] pixels, int w, int h) {
        int len = pixels.length;
        byte[] result = MembandBoost.allocateUninitializedByteArray(len * 3);
        for (int i = 0; i < len; i++) {
            int pixel = pixels[i];
            int idx = i * 3;
            result[idx] = (byte) ((pixel >> 16) & 0xFF);
            result[idx + 1] = (byte) ((pixel >> 8) & 0xFF);
            result[idx + 2] = (byte) (pixel & 0xFF);
        }
        return result;
    }

    public static int[] deserialize(byte[] data, int w, int h) throws IOException {
        int expectedLen = w * h * 3;
        if (data.length != expectedLen) {
            throw new IOException("数据长度不匹配: 期望 " + expectedLen + " 字节, 实际 " + data.length + " 字节");
        }
        int[] pixels = MembandBoost.allocateUninitializedIntArray(w * h);
        for (int i = 0; i < pixels.length; i++) {
            int idx = i * 3;
            int r = data[idx] & 0xFF;
            int g = data[idx + 1] & 0xFF;
            int b = data[idx + 2] & 0xFF;
            pixels[i] = (r << 16) | (g << 8) | b;
        }
        return pixels;
    }

    // ==================== 合并 Zigzag (RGB, 3通道) ====================

    public static byte[] serializeWithZigzag(int[] pixels, int w, int h) {
        int len = pixels.length;
        byte[] result = MembandBoost.allocateUninitializedByteArray(len * 3);
        for (int i = 0; i < len; i++) {
            int pixel = pixels[i];
            byte r = (byte) ((pixel >> 16) & 0xFF);
            byte g = (byte) ((pixel >> 8) & 0xFF);
            byte b = (byte) (pixel & 0xFF);
            byte mr = Zigzag.map(r);
            byte mg = Zigzag.map(g);
            byte mb = Zigzag.map(b);
            int idx = i * 3;
            result[idx] = mr;
            result[idx + 1] = mg;
            result[idx + 2] = mb;
        }
        return result;
    }

    public static int[] deserializeWithZigzag(byte[] data, int w, int h) throws IOException {
        int expectedLen = w * h * 3;
        if (data.length != expectedLen) {
            throw new IOException("数据长度不匹配: 期望 " + expectedLen + " 字节, 实际 " + data.length + " 字节");
        }
        int[] pixels = MembandBoost.allocateUninitializedIntArray(w * h);
        for (int i = 0; i < pixels.length; i++) {
            int idx = i * 3;
            byte r = Zigzag.unmap(data[idx]);
            byte g = Zigzag.unmap(data[idx + 1]);
            byte b = Zigzag.unmap(data[idx + 2]);
            pixels[i] = ((r & 0xFF) << 16) | ((g & 0xFF) << 8) | (b & 0xFF);
        }
        return pixels;
    }

    // ==================== 平面序列化 (RGB, 3通道) ====================

    public static byte[] serializePlannarWithZigzag(int[] pixels, int w, int h) {
        int len = pixels.length;
        byte[] result = MembandBoost.allocateUninitializedByteArray(len * 3);
        int rOffset = 0;
        int gOffset = len;
        int bOffset = len * 2;
        for (int i = 0; i < len; i++) {
            int pixel = pixels[i];
            byte r = (byte) ((pixel >> 16) & 0xFF);
            byte g = (byte) ((pixel >> 8) & 0xFF);
            byte b = (byte) (pixel & 0xFF);
            result[rOffset + i] = Zigzag.map(r);
            result[gOffset + i] = Zigzag.map(g);
            result[bOffset + i] = Zigzag.map(b);
        }
        return result;
    }

    public static int[] deserializePlannarWithZigzag(byte[] data, int w, int h) throws IOException {
        int expectedLen = w * h * 3;
        if (data.length != expectedLen) {
            throw new IOException("数据长度不匹配: 期望 " + expectedLen + " 字节, 实际 " + data.length + " 字节");
        }
        int len = w * h;
        int[] pixels = MembandBoost.allocateUninitializedIntArray(len);
        int rOffset = 0;
        int gOffset = len;
        int bOffset = len * 2;
        for (int i = 0; i < len; i++) {
            byte r = Zigzag.unmap(data[rOffset + i]);
            byte g = Zigzag.unmap(data[gOffset + i]);
            byte b = Zigzag.unmap(data[bOffset + i]);
            pixels[i] = ((r & 0xFF) << 16) | ((g & 0xFF) << 8) | (b & 0xFF);
        }
        return pixels;
    }

    // ==================== RGBA 基础序列化 (4通道) ====================

    /**
     * RGBA 序列化：每个像素 [R, G, B, A] 四个字节（交错存储）
     */
    public static byte[] serializeRGBA(int[] pixels, int w, int h) {
        int len = pixels.length;
        byte[] result = MembandBoost.allocateUninitializedByteArray(len * 4);
        for (int i = 0; i < len; i++) {
            int pixel = pixels[i];
            int idx = i * 4;
            result[idx] = (byte) ((pixel >> 16) & 0xFF);
            result[idx + 1] = (byte) ((pixel >> 8) & 0xFF);
            result[idx + 2] = (byte) (pixel & 0xFF);
            result[idx + 3] = (byte) ((pixel >> 24) & 0xFF);
        }
        return result;
    }

    /**
     * RGBA 反序列化
     */
    public static int[] deserializeRGBA(byte[] data, int w, int h) throws IOException {
        int expectedLen = w * h * 4;
        if (data.length != expectedLen) {
            throw new IOException("数据长度不匹配: 期望 " + expectedLen + " 字节, 实际 " + data.length + " 字节");
        }
        int[] pixels = MembandBoost.allocateUninitializedIntArray(w * h);
        for (int i = 0; i < pixels.length; i++) {
            int idx = i * 4;
            int r = data[idx] & 0xFF;
            int g = data[idx + 1] & 0xFF;
            int b = data[idx + 2] & 0xFF;
            int a = data[idx + 3] & 0xFF;
            pixels[i] = (a << 24) | (r << 16) | (g << 8) | b;
        }
        return pixels;
    }

    // ==================== RGBA + Zigzag (交错) ====================

    public static byte[] serializeRGBAWithZigzag(int[] pixels, int w, int h) {
        int len = pixels.length;
        byte[] result = MembandBoost.allocateUninitializedByteArray(len * 4);
        for (int i = 0; i < len; i++) {
            int pixel = pixels[i];
            byte r = (byte) ((pixel >> 16) & 0xFF);
            byte g = (byte) ((pixel >> 8) & 0xFF);
            byte b = (byte) (pixel & 0xFF);
            byte a = (byte) ((pixel >> 24) & 0xFF);
            byte mr = Zigzag.map(r);
            byte mg = Zigzag.map(g);
            byte mb = Zigzag.map(b);
            byte ma = Zigzag.map(a);  // Alpha 也做 Zigzag
            int idx = i * 4;
            result[idx] = mr;
            result[idx + 1] = mg;
            result[idx + 2] = mb;
            result[idx + 3] = ma;
        }
        return result;
    }

    public static int[] deserializeRGBAWithZigzag(byte[] data, int w, int h) throws IOException {
        int expectedLen = w * h * 4;
        if (data.length != expectedLen) {
            throw new IOException("数据长度不匹配: 期望 " + expectedLen + " 字节, 实际 " + data.length + " 字节");
        }
        int[] pixels = MembandBoost.allocateUninitializedIntArray(w * h);
        for (int i = 0; i < pixels.length; i++) {
            int idx = i * 4;
            byte r = Zigzag.unmap(data[idx]);
            byte g = Zigzag.unmap(data[idx + 1]);
            byte b = Zigzag.unmap(data[idx + 2]);
            byte a = Zigzag.unmap(data[idx + 3]);
            pixels[i] = ((a & 0xFF) << 24) | ((r & 0xFF) << 16) | ((g & 0xFF) << 8) | (b & 0xFF);
        }
        return pixels;
    }

    // ==================== RGBA 平面序列化 (RRRR...GGGG...BBBB...AAAA...) ====================

    public static byte[] serializePlannarRGBAWithZigzag(int[] pixels, int w, int h) {
        int len = pixels.length;
        byte[] result = MembandBoost.allocateUninitializedByteArray(len * 4);
        int rOffset = 0;
        int gOffset = len;
        int bOffset = len * 2;
        int aOffset = len * 3;
        for (int i = 0; i < len; i++) {
            int pixel = pixels[i];
            byte r = (byte) ((pixel >> 16) & 0xFF);
            byte g = (byte) ((pixel >> 8) & 0xFF);
            byte b = (byte) (pixel & 0xFF);
            byte a = (byte) ((pixel >> 24) & 0xFF);
            result[rOffset + i] = Zigzag.map(r);
            result[gOffset + i] = Zigzag.map(g);
            result[bOffset + i] = Zigzag.map(b);
            result[aOffset + i] = Zigzag.map(a);  // Alpha 单独一个平面
        }
        return result;
    }

    public static int[] deserializePlannarRGBAWithZigzag(byte[] data, int w, int h) throws IOException {
        int expectedLen = w * h * 4;
        if (data.length != expectedLen) {
            throw new IOException("数据长度不匹配: 期望 " + expectedLen + " 字节, 实际 " + data.length + " 字节");
        }
        int len = w * h;
        int[] pixels = MembandBoost.allocateUninitializedIntArray(len);
        int rOffset = 0;
        int gOffset = len;
        int bOffset = len * 2;
        int aOffset = len * 3;
        for (int i = 0; i < len; i++) {
            byte r = Zigzag.unmap(data[rOffset + i]);
            byte g = Zigzag.unmap(data[gOffset + i]);
            byte b = Zigzag.unmap(data[bOffset + i]);
            byte a = Zigzag.unmap(data[aOffset + i]);
            pixels[i] = ((a & 0xFF) << 24) | ((r & 0xFF) << 16) | ((g & 0xFF) << 8) | (b & 0xFF);
        }
        return pixels;
    }

    // ==================== 工具方法 ====================

    public static BufferedImage pixelsToBufferedImage(int[] pixels, int w, int h) {
        BufferedImage image = new BufferedImage(w, h, BufferedImage.TYPE_INT_ARGB);
        int[] data = ((DataBufferInt) image.getRaster().getDataBuffer()).getData();
        System.arraycopy(pixels, 0, data, 0, pixels.length);
        return image;
    }
}