package net.northking.cctp.attachment.util;

import com.jcraft.jzlib.Deflater;
import com.jcraft.jzlib.GZIPException;
import com.jcraft.jzlib.Inflater;
import com.jcraft.jzlib.JZlib;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.zip.CRC32;

public class NorthKingIpaPngFix {
    /**
     * 源数据
     */
    private byte[] source;

    /**
     * Png的Trunk信息
     */
    private ArrayList<PNGTrunk> trunks = new ArrayList<>();

    /**
     * 是不是Xcode特殊处理的编码
     */
    private boolean isCgBI = false;

    private PNGIHDRTrunk ihdrTrunk = null;

    private ArrayList<PNGTrunk> idatTrunks = new ArrayList<>();

    private ArrayList<PNGTrunk> headTrunks = new ArrayList<>();

    /**
     * 密封标记。保证对象只能使用一次。
     */
    private boolean isSealed = false;

    public NorthKingIpaPngFix(byte[] source) throws IOException {
        this.source = source;
        readNecessaryChunk();
        if (idatTrunks.isEmpty()) throw new IOException("未能读到有效图片像素数据");
    }

    public NorthKingIpaPngFix(){

    }
    public boolean needConvert() {
        return isCgBI && ihdrTrunk != null;
    }

    /**
     * 转换并获得内容
     *
     * @param formatName 编码格式名称，例如png
     * @return png文件原始数据
     * @throws IOException 发生转换错误
     */
    public byte[] convert(String formatName) throws IOException {
        if (isSealed) throw new IllegalStateException("此对象已经使用过了");
        isSealed = true;

        if (!needConvert()) return source;

        int nMaxInflateBuffer = 4 * (ihdrTrunk.m_nWidth + 1) * ihdrTrunk.m_nHeight;
        byte[] outputBuffer = new byte[nMaxInflateBuffer];
        long inflatedSize = inflate(outputBuffer);

        Deflater deflater = deflate(outputBuffer, (int) inflatedSize, nMaxInflateBuffer);
        PNGTrunk firstDataTrunk = idatTrunks.get(0);

        CRC32 crc32 = new CRC32();
        crc32.update(firstDataTrunk.getName().getBytes());
        crc32.update(deflater.getNextOut(), 0, (int) deflater.getTotalOut());
        long lCRCValue = crc32.getValue();

        firstDataTrunk.m_nData = deflater.getNextOut();
        firstDataTrunk.m_nCRC[0] = (byte) ((lCRCValue & 0xFF000000) >> 24);
        firstDataTrunk.m_nCRC[1] = (byte) ((lCRCValue & 0xFF0000) >> 16);
        firstDataTrunk.m_nCRC[2] = (byte) ((lCRCValue & 0xFF00) >> 8);
        firstDataTrunk.m_nCRC[3] = (byte) (lCRCValue & 0xFF);
        firstDataTrunk.m_nSize = (int) deflater.getTotalOut();
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try {
            byte[] pngHeader = {-119, 80, 78, 71, 13, 10, 26, 10};
            outputStream.write(pngHeader);
            for (int i = 0; i < headTrunks.size() - 1; i++) {
                headTrunks.get(i).writeToStream(outputStream);
            }
            idatTrunks.get(0).writeToStream(outputStream);
            //IEND trunk要在最后
            headTrunks.get(headTrunks.size() - 1).writeToStream(outputStream);
            outputStream.flush();
        } finally {
            outputStream.close();
        }

        source = null;
        trunks = null;
        ihdrTrunk = null;
        idatTrunks = null;
        headTrunks = null;

        byte[] fixedByteArray = outputStream.toByteArray();
        outputBuffer.clone();
        outputStream.close();

        //使用ImageIO解码，得到色彩后修正，再重新编码，也许以后可以优化
        BufferedImage bufferedImage = ImageIO.read(new ByteArrayInputStream(fixedByteArray));
        for (int x = 0; x < bufferedImage.getWidth(); x++) {
            for (int y = 0; y < bufferedImage.getHeight(); y++) {
                int color = bufferedImage.getRGB(x, y);
                int blue = color & 0x0000FF;
                int green = (color & 0x00FF00) >> 8;
                int red = (color & 0xFF0000) >> 16;
                bufferedImage.setRGB(x, y, getIntFromColor(color & 0xFF000000, blue, green, red));
            }
        }
        ByteArrayOutputStream fixColorByteArrayOutputStream = new ByteArrayOutputStream();
        ImageIO.write(bufferedImage, formatName, fixColorByteArrayOutputStream);
        return fixColorByteArrayOutputStream.toByteArray();
    }

    private int getIntFromColor(int alpha, int Red, int Green, int Blue) {
        Red = (Red << 16) & 0x00FF0000;
        Green = (Green << 8) & 0x0000FF00;
        Blue = Blue & 0x000000FF;

        return alpha | Red | Green | Blue; //0xFF000000 for 100% Alpha. Bitwise OR everything together.
    }

    /**
     * 读取CgBI和IHDR信息
     */
    private void readNecessaryChunk() throws IOException {
        try (DataInputStream input = new DataInputStream(new ByteArrayInputStream(source))) {
            byte[] nPNGHeader = new byte[8];
            input.readFully(nPNGHeader);

            if ((nPNGHeader[0] == -119) && (nPNGHeader[1] == 0x50) && (nPNGHeader[2] == 0x4e) && (nPNGHeader[3] == 0x47)
                    && (nPNGHeader[4] == 0x0d) && (nPNGHeader[5] == 0x0a) && (nPNGHeader[6] == 0x1a) && (nPNGHeader[7] == 0x0a)) {

                PNGTrunk trunk;
                do {
                    trunk = PNGTrunk.generateTrunk(input);
                    trunks.add(trunk);
                    boolean isHeadTrunk = true;
                    if (trunk.getName().equalsIgnoreCase("CgBI")) {
                        isCgBI = true;
                        isHeadTrunk = false;
                    }
                    if (trunk.getName().equalsIgnoreCase("IHDR")) {
                        ihdrTrunk = (PNGIHDRTrunk) trunk;
                    }
                    if (trunk.getName().equalsIgnoreCase("IDAT")) {
                        idatTrunks.add(trunk);
                        isHeadTrunk = false;
                    }
                    if (isHeadTrunk) {
                        headTrunks.add(trunk);
                    }
                }
                while (!trunk.getName().equalsIgnoreCase("IEND"));
            }
        }
    }

    private long inflate(byte[] conversionBuffer) throws GZIPException {
        Inflater inflater = new Inflater(-15);

        for (PNGTrunk dataTrunk : idatTrunks) {
            inflater.setInput(dataTrunk.getData(), true);
        }

        inflater.setOutput(conversionBuffer);

        int result;
        try {
            result = inflater.inflate(JZlib.Z_NO_FLUSH);
            checkResultStatus(result);
        } finally {
            inflater.inflateEnd();
        }

        return inflater.getTotalOut();
    }

    private void checkResultStatus(int nResult) throws GZIPException {
        switch (nResult) {
            case JZlib.Z_OK:
            case JZlib.Z_STREAM_END:
                break;

            case JZlib.Z_NEED_DICT:
                throw new GZIPException("Z_NEED_DICT - " + nResult);
            case JZlib.Z_DATA_ERROR:
                throw new GZIPException("Z_DATA_ERROR - " + nResult);
            case JZlib.Z_MEM_ERROR:
                throw new GZIPException("Z_MEM_ERROR - " + nResult);
            case JZlib.Z_STREAM_ERROR:
                throw new GZIPException("Z_STREAM_ERROR - " + nResult);
            case JZlib.Z_BUF_ERROR:
                throw new GZIPException("Z_BUF_ERROR - " + nResult);
            default:
                throw new GZIPException("inflater error: " + nResult);
        }
    }

    private Deflater deflate(byte[] buffer, int length, int nMaxInflateBuffer) throws GZIPException {
        Deflater deflater = new Deflater();
        deflater.setInput(buffer, 0, length, false);

        int nMaxDeflateBuffer = nMaxInflateBuffer + 1024;
        byte[] deBuffer = new byte[nMaxDeflateBuffer];
        deflater.setOutput(deBuffer);

        deflater.deflateInit(JZlib.Z_BEST_COMPRESSION);
        int nResult = deflater.deflate(JZlib.Z_FINISH);
        checkResultStatus(nResult);

        if (deflater.getTotalOut() > nMaxDeflateBuffer) {
            throw new GZIPException("deflater output buffer was too small");
        }

        return deflater;
    }
}
