package org.kne.codec.kif;

import io.airlift.compress.Compressor;
import io.airlift.compress.Decompressor;
import io.airlift.compress.thirdparty.JdkDeflateCompressor;
import io.airlift.compress.thirdparty.JdkInflateDecompressor;
import io.airlift.compress.thirdparty.ZstdJniCompressor;
import io.airlift.compress.zstd.ZstdCompressor;
import io.airlift.compress.zstd.ZstdDecompressor;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.zip.Deflater;
import java.util.zip.Inflater;

import org.kne.membandboost.MembandBoost;

/**
 * 压缩/解压缩工具
 * 支持 Deflate（ZLIB）和 Zstandard 两种算法
 */
public class Compressors {

    // ==================== Zstandard (Zstd) ====================
    private static final ThreadLocal<ZstdCompressor> ZSTD_COMPRESSOR = 
    	    ThreadLocal.withInitial(ZstdCompressor::new);
    	    
    	private static final ThreadLocal<ZstdDecompressor> ZSTD_DECOMPRESSOR = 
    	    ThreadLocal.withInitial(ZstdDecompressor::new);

        // ==================== Deflate (ZLIB) ====================

        private static final ThreadLocal<JdkDeflateCompressor> DEFLATE_COMPRESSOR = 
        	    ThreadLocal.withInitial(JdkDeflateCompressor::new);
        	    
        	private static final ThreadLocal<JdkInflateDecompressor> DEFLATE_DECOMPRESSOR = 
        	    ThreadLocal.withInitial(JdkInflateDecompressor::new);
    /**
     * 压缩
     */
    public static byte[] compress(Compressor comp,byte[] input) {
        int maxLen = comp.maxCompressedLength(input.length);
        byte[] compressed = MembandBoost.allocateUninitializedByteArray(maxLen);
        int compressedSize = comp.compress(
            input, 0, input.length,
            compressed, 0,
            maxLen
        );
        //System.out.println("in:"+input.length+" out:"+compressedSize);
        byte[] copy = MembandBoost.allocateUninitializedByteArray(compressedSize);
		System.arraycopy(compressed, 0, copy, 0,
		                 Math.min(compressed.length, compressedSize));
		return copy;
    }

    /**
     * 解压（已知原始长度）
     */
    public static byte[] decompress(Decompressor decomp,byte[] input, int expectedLength) throws IOException {
    	try {
    		
        byte[] output = MembandBoost.allocateUninitializedByteArray(expectedLength);
        int decompressedSize = decomp.decompress(
            input, 0, input.length,
            output, 0, expectedLength
        );
        if (decompressedSize != expectedLength) {
            throw new IOException("解压长度不匹配: 期望 " + expectedLength + ", 实际 " + decompressedSize);
        }
        return output;
    	}catch(Exception e) {
    		e.printStackTrace();
        System.out.println("in:"+input.length+" expected:"+expectedLength);
        throw e;
    	}
    }
    

    public static byte[] compressZstd(byte[] input) {
    	return compress(ZSTD_COMPRESSOR.get(), input);
    }
    public static byte[] decompressZstd(byte[] input, int expectedLength) throws IOException {
    	return decompress(ZSTD_DECOMPRESSOR.get(), input, expectedLength);
    }

    public static byte[] compressDeflate(byte[] input) {
    	return compress(DEFLATE_COMPRESSOR.get(), input);
    }
    public static byte[] decompressDeflate(byte[] input, int expectedLength) throws IOException {
    	return decompress(DEFLATE_DECOMPRESSOR.get(), input, expectedLength);
    }
}