package com.cscn;

import javacard.framework.JCSystem;

import static com.cscn.Zuc256Core.zuc256GenerateKeystream;
import static com.cscn.Zuc256Core.zuc256GenerateKeyword;
import static com.cscn.Zuc256Util.getU32;
import static com.cscn.Zuc256Util.putU32;
import static com.cscn.Zuc256Util.xor32;


/**
 * 加密上下文类
 */
public final class Zuc256EncryptCtx {
    Zuc256State state;
    byte[] buf;
    short buflen;

    public Zuc256EncryptCtx(Zuc256State state, byte[] buf){
        this.state = state;
        this.buf = buf;
    }

    public Zuc256EncryptCtx(Zuc256State state){
        this.state = state;
        this.buf = new byte[4];
    }

    public Zuc256EncryptCtx(){
        this.state = new Zuc256State();
        this.buf = new byte[4];
    }

    // 初始化加密上下文
    public void init(byte[] key32, byte[] iv) {
//        Arrays.fill(this.buf, (byte) 0);
        for (short i = 0; i < (short)this.buf.length; i++) {
            this.buf[i] = (byte)0;
        }
        this.buflen = 0;
        Zuc256Core.initState(this.state, key32, iv);
    }

    // 分阶段处理加密数据
    public void update(byte[] in, short inlen, byte[] out) {
        if (in == null || out == null || inlen == 0) return;

        short inPos  = 0; // 输入偏移
        short outPos = 0; // 输出偏移

        // 处理缓冲区中剩余的非4字节数据
        if (this.buflen > 0) {
//            int need = 4 - this.buflen;
            short need = (short)(4 - this.buflen);
//            int copy = Math.min(inlen, need);
            short copy = (short)((inlen < need) ? inlen : need);

            System.arraycopy(in, 0, this.buf, this.buflen, copy);
            this.buflen += copy;

            // 调整输入指针和长度
//            byte[] newIn = new byte[inlen - copy];
//            if (inlen - copy > 0) {
//                System.arraycopy(in, copy, newIn, 0, inlen - copy);
//            }
//            in = newIn;
//            inlen -= copy;
            // 推进输入指针与剩余长度
            inPos  += copy;
            inlen  -= copy;

            // 缓冲区已满，处理一个完整的4字节块
            if (this.buflen == 4) {
//                int keystream = zuc256GenerateKeyword(this.state);
                short[] ks = JCSystem.makeTransientShortArray((short)2, JCSystem.CLEAR_ON_DESELECT);
                zuc256GenerateKeyword(this.state, ks);  // ks[0]=lo, ks[1]=hi

//                int plain = getU32(this.buf, 0);
                // 取出 4 字节明文 → plain[0]=lo, plain[1]=hi
                short[] plain = JCSystem.makeTransientShortArray((short)2, JCSystem.CLEAR_ON_DESELECT);
                getU32(this.buf, (short)0, plain);

//                putU32(out, 0, plain ^ keystream);
                // plain ^ ks → res
                short[] res = JCSystem.makeTransientShortArray((short)2, JCSystem.CLEAR_ON_DESELECT);
                xor32(plain[0], plain[1], ks[0], ks[1], res);
                // 写回 out 的前4字节
                putU32(out, outPos, res[0], res[1]);

                this.buflen = 0;
//                Arrays.fill(this.buf, (byte) 0);
                for (short i = 0; i < (short)this.buf.length; i++) {
                    this.buf[i] = (byte)0;
                }

                // 调整输出指针
//                byte[] newOut = new byte[out.length - 4];
//                if (out.length - 4 > 0) {
//                    System.arraycopy(out, 4, newOut, 0, out.length - 4);
//                }
//                out = newOut;
                // 这里C实现就是直接指针+4的。JavaSE实现搞这个new干嘛。。
                outPos += 4;
            }
        }

        // 处理完整的4字节块
//        int fullBlocks = inlen / 4;
        short fullBlocks = (short) (inlen / 4);
        if (fullBlocks > 0) {
//            int[] keystream = new int[fullBlocks];
            short[] ks_hi = JCSystem.makeTransientShortArray(fullBlocks, JCSystem.CLEAR_ON_DESELECT);
            short[] ks_lo = JCSystem.makeTransientShortArray(fullBlocks, JCSystem.CLEAR_ON_DESELECT);

//            zuc256GenerateKeystream(this.state, fullBlocks, keystream);
            zuc256GenerateKeystream(this.state, fullBlocks, ks_hi, ks_lo);

            // 临时：装一个32位字
            short[] word = JCSystem.makeTransientShortArray((short)2, JCSystem.CLEAR_ON_DESELECT);

            // 逐块异或加密
            for (short i = 0; i < fullBlocks; i++) {
//                int plain = getU32(in, i * 4);
                short off = (short) (i << 2); // i*4
                // 读明文
                getU32(in, (short)(inPos+off), word);   // word[0]=lo, word[1]=hi

//                putU32(out, i * 4, plain ^ keystream[i]);
                // XOR keystream
                word[0] = (short)(word[0] ^ ks_lo[i]);
                word[1] = (short)(word[1] ^ ks_hi[i]);
                // 写密文
                putU32(out, (short) (outPos+off), word[0], word[1]);
            }

            // 调整输入指针和长度
//            int processed = fullBlocks * 4;
            short processed = (short)(fullBlocks * 4);

//            byte[] newIn = new byte[inlen - processed];
//            if (inlen - processed > 0) {
//                System.arraycopy(in, processed, newIn, 0, inlen - processed);
//            }
//            in = newIn;
//            inlen -= processed;
            // 推进输入/输出指针与剩余长度
            inPos  += processed;
            inlen  -= processed;
            outPos += processed;
        }

        // 缓存剩余不足4字节的数据
        if (inlen > 0) {
            System.arraycopy(in, 0, this.buf, 0, inlen);
            this.buflen = inlen;
        }
    }

    // 完成加密处理
    public void finish(byte[] out) {
        if (out == null) return;

        // 处理缓冲区中剩余的不足4字节数据
        if (this.buflen > 0) {
//            int keystream = zuc256GenerateKeyword(this.state);
            // 生成一个 32-bit 密钥字：ks[0]=lo16, ks[1]=hi16
            short[] ks = JCSystem.makeTransientShortArray((short)2, JCSystem.CLEAR_ON_DESELECT);
            zuc256GenerateKeyword(this.state, ks);

//            byte[] keystreamBytes = new byte[4];
//            putU32(keystreamBytes, 0, keystream);
            byte[] keystreamBytes = new byte[4];
            putU32(keystreamBytes, (short)0, ks[0], ks[1]);

            // 逐字节异或
            for (short i = 0; i < this.buflen; i++) {
                out[i] = (byte) (this.buf[i] ^ keystreamBytes[i]);
            }
        }

        // 清理上下文
//        Arrays.fill(this.buf, (byte) 0);
        for (short i = 0; i < (short)this.buf.length; i++) {
            this.buf[i] = (byte)0;
        }

        this.buflen = 0;
//        Arrays.fill(this.state.LFSR, 0);
        // LFSR 全部清零（高低位数组各 16 个元素）
        for (short i = 0; i < 16; i++) {
            this.state.LFSR_lo[i] = 0;
            this.state.LFSR_hi[i] = 0;
        }

//        this.state.R1 = 0;
//        this.state.R2 = 0;
// R1、R2 清零
        this.state.R1_lo = 0;
        this.state.R1_hi = 0;
        this.state.R2_lo = 0;
        this.state.R2_hi = 0;
    }
}