package api;

import com.sunyard.chsm.enums.AlgMode;
import com.sunyard.chsm.enums.Padding;
import com.sunyard.chsm.param.*;
import com.sunyard.chsm.sdf.util.LangUtils;
import com.sunyard.chsm.utils.CodecUtils;
import com.sunyard.chsm.utils.gm.BCSM3Utils;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.springframework.util.Assert;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.Collections;

/**
 * @author liulu
 * @since 2024/12/27
 */
@Slf4j
public class BigDataTest extends BaseTest {

    private static final byte[] plain = new byte[1024 * 7];
    private final static byte[] iv = "ghwikdhj1234713v".getBytes();
    private static Long keyId;

    @BeforeAll
    public static void beforeAll() {
        Arrays.fill(plain, (byte) 0x04);
        keyId = execute("/key/gen", KeyCreateReq.builder().keyTemplateCode(keyTemplate).genNumber(1).build(), Long.class);
        Assertions.assertTrue(keyId > 0);
    }

    @AfterAll
    public static void afterAll() {
        KeyManageReq keyManageReq = new KeyManageReq();
        keyManageReq.setIds(Collections.singletonList(keyId));
        execute("/key/disable", keyManageReq, Void.class);
        execute("/key/destroy", keyManageReq, Void.class);
    }

    @Test
    public void testMultiSm4() {

        byte[] bigdata = new byte[1024*7*1024];
        Arrays.fill(bigdata, (byte) 0x07);

        SymEncryptResp resp = multiPackageSymEnc(bigdata);
        byte[] cipherData = CodecUtils.decodeBase64(resp.getCipherData());
        log.info("SymEncryptResp: {}, {}", cipherData.length, resp);

        byte[] plain = multiPackageSymDec(resp.getKeyIndex(), cipherData);
        log.info("multiPackageSymDec: {}", plain.length);

        Assertions.assertArrayEquals(bigdata, plain);
    }


    @Test
    public void testMultiHash() {

        byte[] bigdata = new byte[1024 * 7 * 128];
        Arrays.fill(bigdata, (byte) 0x07);

        byte[] hash = multiPackageHash(bigdata);
        log.info("sdf hash: {}", CodecUtils.encodeHex(hash));
        byte[] bcHash = BCSM3Utils.hash(bigdata);
        log.info("bc hash: {}", CodecUtils.encodeHex(bcHash));
        Assertions.assertArrayEquals(hash, bcHash);
    }


    public static byte[] multiPackageHash(byte[] data) {
        return multiPackageHash(new ByteArrayInputStream(data));
    }

    public static byte[] multiPackageHash(InputStream is) {
        HashResp hashResp = execute("/multi/package/hash/init", new HashReq(), HashResp.class);
        Assert.notNull(hashResp, "HashInit failed");
        log.info("hash: {}", hashResp);

        HashReq hashReq = new HashReq();
        hashReq.setHandle(hashResp.getHandle());
        int len = 1024 * 7;
        byte[] buffer = new byte[len];
        while (true) {
            int bytesRead;
            try {
                if ((bytesRead = is.read(buffer)) == -1) break;
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            if (bytesRead < len) {
                buffer = Arrays.copyOf(buffer, bytesRead);
            }
            hashReq.setPlainData(CodecUtils.encodeBase64(buffer));
            hashResp = execute("/multi/package/hash/update", hashReq, HashResp.class);
        }

        hashResp = execute("/multi/package/hash/finish", hashReq, HashResp.class);
        Assert.notNull(hashResp, "HashFinish failed");
        log.info("hash: {}", hashResp);
        return CodecUtils.decodeBase64(hashResp.getHash());
    }

    static int len = 1024*7;
    public static SymEncryptResp multiPackageSymEnc(byte[] array) {
        return multiPackageSymEnc(new ByteArrayInputStream(array));
    }

    public static SymEncryptResp multiPackageSymEnc(InputStream is) {
        SymEncryptReq symEncryptReq = new SymEncryptReq();
        symEncryptReq.setKeyId(keyId);
        symEncryptReq.setPadding(Padding.NOPadding);
        symEncryptReq.setIv(CodecUtils.encodeBase64(iv));
        symEncryptReq.setMode(AlgMode.CBC);

        SymEncryptResp resp = new SymEncryptResp();
        byte[] buffer = new byte[len];
        try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
            while (true) {
                int bytesRead;
                if ((bytesRead = is.read(buffer)) == -1) break;

                if (bytesRead < len) {
                    symEncryptReq.setPadding(Padding.PCKS7Padding);
                    buffer = Arrays.copyOf(buffer, bytesRead);
                }
                int available = is.available();
                if (available < 16) {
                    byte[] bytes = new byte[available];
                    is.read(bytes);
                    buffer = LangUtils.merge(buffer, bytes);
                    symEncryptReq.setPadding(Padding.PCKS7Padding);
                }

                symEncryptReq.setPlainData(CodecUtils.encodeBase64(buffer));
                SymEncryptResp tmp = execute("/sym/encrypt", symEncryptReq, SymEncryptResp.class);
                os.write(CodecUtils.decodeBase64(tmp.getCipherData()));
                resp.setKeyId(tmp.getKeyId());
                resp.setKeyIndex(tmp.getKeyIndex());
            }
            resp.setCipherData(CodecUtils.encodeBase64(os.toByteArray()));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return resp;
    }

    public static byte[] multiPackageSymDec(String keyIdx, byte[] array) {
        return multiPackageSymDec(keyIdx, new ByteArrayInputStream(array));
    }

    public static byte[] multiPackageSymDec(String keyIdx, InputStream is) {
        SymDecryptReq decryptReq = new SymDecryptReq();
        decryptReq.setKeyId(keyId);
        decryptReq.setKeyIndex(keyIdx);
        decryptReq.setPadding(Padding.NOPadding);
        decryptReq.setIv(CodecUtils.encodeBase64(iv));
        decryptReq.setMode(AlgMode.CBC);

        byte[] buffer = new byte[len];
        try (ByteArrayOutputStream os = new ByteArrayOutputStream()) {
            while (true) {
                int bytesRead;
                if ((bytesRead = is.read(buffer)) == -1) break;

                if (bytesRead < len) {
                    decryptReq.setPadding(Padding.PCKS7Padding);
                    buffer = Arrays.copyOf(buffer, bytesRead);
                }
                int available = is.available();
                if (available <= 16) {
                    decryptReq.setPadding(Padding.PCKS7Padding);
                    byte[] bytes = new byte[available];
                    is.read(bytes);
                    buffer = LangUtils.merge(buffer, bytes);
                }
                decryptReq.setCipherData(CodecUtils.encodeBase64(buffer));
                SymDecryptResp resp = execute("/sym/decrypt", decryptReq, SymDecryptResp.class);
                os.write(CodecUtils.decodeBase64(resp.getPlainData()));
            }
            return os.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }


}
