package org.kne.codec.kif;

import java.awt.image.BufferedImage;
import java.awt.image.DataBufferInt;
import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

import jdk.internal.foreign.ArenaImpl;
import org.kne.membandboost.MembandBoost;

public class IntImageBlock {
    private MemorySegment image; // int[] 的 MemorySegment 封装
    private int width;
    private int height;
    private int posx;
    private int posy;

    public IntImageBlock(MemorySegment image, int width, int height, int posx, int posy) {
        this.image = image;
        this.width = width;
        this.height = height;
        this.posx = posx;
        this.posy = posy;
    }

    public IntImageBlock(MemorySegment image, int width, int height) {
        this.image = image;
        this.width = width;
        this.height = height;
        this.posx = 0;
        this.posy = 0;
    }

    public IntImageBlock(int width2, int height2, int posx2, int posy2) {
        this(((ArenaImpl) Arena.ofAuto()).allocateNoInit((width2 * height2) << 2, 4), width2, height2, posx2, posy2);
    }

    public List<BlockCoord> getBlockCoords(int blockWidth, int blockHeight) {
        return BlockCoord.calculateBlockCoords(width, height, blockWidth, blockHeight);
    }

    public IntImageBlock createBlock(BlockCoord coord) {
        return IntImageBlock.fromImage(image, width, height, coord.getPosX(), coord.getPosY(), coord.getWidth(),
                coord.getHeight());
    }

    public static IntImageBlock createBlock(BufferedImage image, BlockCoord coord) {
        return IntImageBlock.fromImage(image, image.getWidth(), image.getHeight(), coord.getPosX(), coord.getPosY(),
                coord.getWidth(), coord.getHeight());
    }

    // ==================== 从 MemorySegment 提取块 ====================

    public static IntImageBlock fromImage(MemorySegment src, int srcw, int srch, int posx, int posy, int blockw,
            int blockh) {
        if (src == null) {
            throw new IllegalArgumentException("源图像不能为空");
        }
        if (posx < 0 || posy < 0 || blockw <= 0 || blockh <= 0) {
            throw new IllegalArgumentException("坐标和尺寸必须为正数");
        }
        if (posx + blockw > srcw || posy + blockh > srch) {
            throw new IllegalArgumentException("块区域超出源图像边界");
        }

        MemorySegment blockPixels = ((ArenaImpl) Arena.ofAuto()).allocateNoInit((blockw * blockh) << 2, 4);

        // 逐行拷贝（使用 MemorySegment.copy）
        for (int row = 0; row < blockh; row++) {
            long srcOffset = ((long) (posy + row) * srcw + posx) << 2;
            long dstOffset = (long) row * blockw << 2;
            MemorySegment.copy(src, srcOffset, blockPixels, dstOffset, (long) blockw << 2);
        }

        return new IntImageBlock(blockPixels, blockw, blockh, posx, posy);
    }

    public static IntImageBlock fromImage(BufferedImage src, int srcw, int srch, int posx, int posy, int blockw,
            int blockh) {
        if (src == null) {
            throw new IllegalArgumentException("源图像不能为空");
        }
        if (posx < 0 || posy < 0 || blockw <= 0 || blockh <= 0) {
            throw new IllegalArgumentException("坐标和尺寸必须为正数");
        }
        if (posx + blockw > srcw || posy + blockh > srch) {
            throw new IllegalArgumentException("块区域超出源图像边界");
        }

        int type = src.getType();
        MemorySegment blockPixels = ((ArenaImpl) Arena.ofAuto()).allocateNoInit((long) blockw * blockh << 2, 4);

        // 情况1：TYPE_INT_ARGB 或 TYPE_INT_RGB → 零拷贝路径
        if (type == BufferedImage.TYPE_INT_ARGB || type == BufferedImage.TYPE_INT_RGB) {
            DataBufferInt buffer = (DataBufferInt) src.getRaster().getDataBuffer();
            int[] pixels = buffer.getData();
            for (int row = 0; row < blockh; row++) {
                int srcIdx = (posy + row) * srcw + posx;
                MemorySegment.copy(pixels, srcIdx, blockPixels, ValueLayout.JAVA_INT, (long) row * blockw << 2, blockw);
            }
            return new IntImageBlock(blockPixels, blockw, blockh, posx, posy);
        }

        // 情况2：其他格式 → 逐行读取，直接写入 MemorySegment
        else {
            int[] rowBuffer = new int[blockw];
            for (int row = 0; row < blockh; row++) {
                // 一次性读取一行（1 次 JNI 调用）
                src.getRGB(posx, posy + row, blockw, 1, rowBuffer, 0, blockw);
                // 直接写入 MemorySegment（1 次 native 拷贝）
                MemorySegment.copy(rowBuffer, 0, blockPixels, ValueLayout.JAVA_INT, (long) row * blockw << 2, blockw);
            }
            return new IntImageBlock(blockPixels, blockw, blockh, posx, posy);
        }
    }

    // ==================== 写入目标图像 ====================

    public void toImage(MemorySegment dst, int dstw, int dsth) {
        if (dst == null) {
            throw new IllegalArgumentException("目标图像不能为空");
        }
        if (posx < 0 || posy < 0) {
            throw new IllegalArgumentException("块位置不能为负数");
        }
        if (posx + width > dstw || posy + height > dsth) {
            throw new IllegalArgumentException("块区域超出目标图像边界");
        }

        // 逐行拷贝到目标图像
        for (int row = 0; row < height; row++) {
            long srcOffset = (long) row * width << 2;
            long dstOffset = ((long) (posy + row) * dstw + posx) << 2;
            MemorySegment.copy(image, srcOffset, dst, dstOffset, (long) width << 2);
        }
    }

    public void toImage(MemorySegment dst, int dstw, int dsth, int dstPosX, int dstPosY) {
        if (dst == null) {
            throw new IllegalArgumentException("目标图像不能为空");
        }
        if (dstPosX < 0 || dstPosY < 0) {
            throw new IllegalArgumentException("目标位置不能为负数");
        }
        if (dstPosX + width > dstw || dstPosY + height > dsth) {
            throw new IllegalArgumentException("块区域超出目标图像边界");
        }

        for (int row = 0; row < height; row++) {
            long srcOffset = (long) row * width << 2;
            long dstOffset = ((long) (dstPosY + row) * dstw + dstPosX) << 2;
            MemorySegment.copy(image, srcOffset, dst, dstOffset, (long) width << 2);
        }
    }

    // ==================== 转换为 int[]（兼容旧代码） ====================

    public int[] toIntArray() {
        int[] array = new int[width * height];
        MemorySegment.copy(image, ValueLayout.JAVA_INT, 0, array, 0, width * height);
        return array;
    }

    // ==================== Getter/Setter ====================

    public MemorySegment getImage() {
        return image;
    }

    public void setImage(MemorySegment image) {
        this.image = image;
    }

    public int getWidth() {
        return width;
    }

    public int getHeight() {
        return height;
    }

    public int getPosx() {
        return posx;
    }

    public int getPosy() {
        return posy;
    }

    @Override
    public String toString() {
        return "IntImageBlock[" + width + "x" + height + ",(" + posx + "," + posy + ")]";
    }
}