package org.kne.codec.kif;

import javax.imageio.IIOImage;
import javax.imageio.ImageTypeSpecifier;
import javax.imageio.ImageWriteParam;
import javax.imageio.ImageWriter;
import javax.imageio.metadata.IIOMetadata;
import javax.imageio.spi.ImageWriterSpi;
import javax.imageio.stream.ImageOutputStream;

import org.kne.debug.TimeDebugger;

import java.awt.AlphaComposite;
import java.awt.Graphics;
import java.awt.Graphics2D;
import java.awt.Image;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferInt;
import java.awt.image.RenderedImage;
import java.io.IOException;
import java.lang.foreign.MemorySegment;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.function.Supplier;
import java.util.zip.Deflater;

/**
 * KIF格式的实际编码器。
 * 这个类负责将BufferedImage编码为.kif文件。
 */
public class KIFImageWriter extends ImageWriter {
	private static final int BLOCK_SIZE=512;
    private ImageOutputStream outputStream;
   private static ModeStatisticsCollector modeCollector=new ModeStatisticsCollector();
   
    public static ModeStatisticsCollector getModeCollector() {
    	return modeCollector;
    }

	protected KIFImageWriter(ImageWriterSpi originatingProvider) {
        super(originatingProvider);
    }

    @Override
    public void setOutput(Object output) {
        super.setOutput(output);
        if (output instanceof ImageOutputStream) {
            this.outputStream = (ImageOutputStream) output;
        } else {
            this.outputStream = null;
        }
    }

    @Override
    public IIOMetadata getDefaultStreamMetadata(ImageWriteParam param) {
        return null;
    }

    @Override
    public IIOMetadata getDefaultImageMetadata(ImageTypeSpecifier imageType, ImageWriteParam param) {
        return null;
    }

    @Override
    public IIOMetadata convertStreamMetadata(IIOMetadata inData, ImageWriteParam param) {
        return null;
    }

    @Override
    public IIOMetadata convertImageMetadata(IIOMetadata inData, ImageTypeSpecifier imageType, ImageWriteParam param) {
        return null;
    }

    @Override
    public void write(IIOMetadata streamMetadata, IIOImage image, ImageWriteParam param) throws IOException {
        if (outputStream == null) {
            throw new IllegalStateException("Output not set");
        }

        // 检查是否支持此图像类型
        if (image.getRenderedImage() == null) {
            throw new IllegalArgumentException("Image contains no RenderedImage");
        }

        RenderedImage renderedImage = image.getRenderedImage();
        int width = renderedImage.getWidth();
        int height = renderedImage.getHeight();

        
        TimeDebugger tdb=new TimeDebugger("写入时间",false);
        // 方法1：如果 RenderedImage 是 BufferedImage，可以直接操作
        BufferedImage bi = null;
        if (renderedImage instanceof BufferedImage) {
            bi = (BufferedImage) renderedImage;
        } else {
            // 如果不是 BufferedImage，创建一个副本（简化处理）
            bi = BufferedImageMembandBoost.createUninitializedBufferedImage(width, height) ;
           Graphics2D g=  (Graphics2D) bi.getGraphics();
           g.setComposite(AlphaComposite.Src);
           g .drawImage((Image) renderedImage, 0, 0, null);
        }
  /*
     // 在 KIFImageWriter.write() 中
        BufferedImage original = (BufferedImage) image.getRenderedImage();

        // 统一转换为 TYPE_INT_ARGB
        BufferedImage rgbImage =null;
        if(original.getType()!=BufferedImage.TYPE_INT_ARGB) {
      rgbImage  = BufferedImageMembandBoost.createUninitializedBufferedImage( original.getWidth(),  original.getHeight()  );
        Graphics2D g = rgbImage.createGraphics();
        g.setComposite(AlphaComposite.Src);
        g.drawImage(original, 0, 0, null);
        g.dispose();
        }else {
        	rgbImage=original;
        }*/
        tdb.mark("格式转换");
        List<Future<KIFCompressedImageBlock>>encoding=new ArrayList<>(); 
        List<BlockCoord>bcd=BlockCoord.calculateBlockCoords(width,height, BLOCK_SIZE, BLOCK_SIZE);
        KIFCodec.backPressure();
        for(BlockCoord crd:bcd) {
        	final BufferedImage bi2=bi;
    	encoding.add( KIFCodec.encodeAsyncMultimode(()->{
    		return IntImageBlock.createBlock(bi2,crd);
    	},KIFCompressMode.getNaturalImageModes()));
        }
        tdb.mark("分块提交");
       

        // ===== 写入文件头 =====
        // 魔数 (3 bytes) "KIF"
        outputStream.write(0x4B);  // 'K'
        outputStream.write(0x49);  // 'I'
        outputStream.write(0x46);  // 'F'
        
        // 版本号 (1 byte)
        outputStream.writeByte(1);
        
        // 图像类型 (1 byte): 0=灰度, 1=RGB, 2=RGBA
        int imageType = bi.getType();
     //   System.out.println("类型："+imageType);
        int typeFlag = 0;
        if (imageType == BufferedImage.TYPE_INT_RGB || 
            imageType == BufferedImage.TYPE_3BYTE_BGR) {
            typeFlag = 1;
        } else if (imageType == BufferedImage.TYPE_INT_ARGB || 
                   imageType == BufferedImage.TYPE_4BYTE_ABGR) {
            typeFlag = 2;
        }
        outputStream.writeByte(typeFlag);
        
        // 宽度 (4 bytes, 大端序)
        outputStream.writeInt(width);
        
        // 高度 (4 bytes, 大端序)
        outputStream.writeInt(height);

        // ===== 写入压缩的图像数据 =====
       
        outputStream.writeInt(encoding.size()); 
        tdb.mark("写入头");
        while(true) {
        	for (Iterator<Future<KIFCompressedImageBlock>> iterator = encoding.iterator(); iterator.hasNext();) {
        		Future<KIFCompressedImageBlock> v =  iterator.next();
        		try {
            		if(v.isDone()) {
            		KIFCompressedImageBlock blk=v.get();
    				blk.writeToStream(outputStream);
    				modeCollector.addRecord(blk.getPredictMode());
    				iterator.remove();
            		}
    			} catch (InterruptedException e) {
    				e.printStackTrace();
    			} catch (ExecutionException e) {
    				ExceptionTool.throwIOException(e);
    			}
			}
        	if(encoding.isEmpty()) {
        		break;
        	}
        	try {
				Thread.sleep(1);
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
        }
        tdb.mark("等待编码完成");
        // 注意：写入完成后记得 flush
        outputStream.flush();
        tdb.mark("写入文件");
        tdb.print();
    }

    /**
     * 将 BufferedImage 转换为平面像素数组（逐行，每个通道独立）
     * 这是一个辅助方法，实际实现时可根据需要调整
     */
    private byte[][] getPixelBands(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();
        int bands = image.getSampleModel().getNumBands();
        
        byte[][] bandData = new byte[bands][width * height];
        
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int idx = y * width + x;
                int rgb = image.getRGB(x, y);
                // 根据图像类型提取各个通道
                // 简化：假设为 RGB
                if (bands >= 3) {
                    bandData[0][idx] = (byte) ((rgb >> 16) & 0xFF); // R
                    bandData[1][idx] = (byte) ((rgb >> 8) & 0xFF);  // G
                    bandData[2][idx] = (byte) (rgb & 0xFF);         // B
                } else if (bands == 1) {
                    bandData[0][idx] = (byte) (rgb & 0xFF);         // 灰度
                }
            }
        }
        return bandData;
    }

    @Override
    public void dispose() {
    	super.dispose();
    }
}