package com.cscn.zuc256;

import com.cscn.Zuc256Core;
import com.cscn.Zuc256State;

import java.util.Arrays;


/**
 * 加密上下文类
 */
public final class Zuc256EncryptCtx {
    Zuc256State state;
    byte[] buf;
    int 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);
        this.buflen = 0;
        Zuc256Core.initState(this.state, key32, iv);
    }

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

        // 处理缓冲区中剩余的非4字节数据
        if (this.buflen > 0) {
            int need = 4 - this.buflen;
            int copy = Math.min(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;

            // 缓冲区已满，处理一个完整的4字节块
            if (this.buflen == 4) {
                int keystream = zuc256GenerateKeyword(this.state);
                int plain = getU32(this.buf, 0);
                putU32(out, 0, plain ^ keystream);

                this.buflen = 0;
                Arrays.fill(this.buf, (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;
            }
        }

        // 处理完整的4字节块
        int fullBlocks = inlen / 4;
        if (fullBlocks > 0) {
            int[] keystream = new int[fullBlocks];
            zuc256GenerateKeystream(this.state, fullBlocks, keystream);

            // 逐块异或加密
            for (int i = 0; i < fullBlocks; i++) {
                int plain = getU32(in, i * 4);
                putU32(out, i * 4, plain ^ keystream[i]);
            }

            // 调整输入指针和长度
            int processed = fullBlocks * 4;
            byte[] newIn = new byte[inlen - processed];
            if (inlen - processed > 0) {
                System.arraycopy(in, processed, newIn, 0, inlen - processed);
            }
            in = newIn;
            inlen -= processed;
        }

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

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

        // 处理缓冲区中剩余的不足4字节数据
        if (this.buflen > 0) {
            int keystream = zuc256GenerateKeyword(this.state);
            byte[] keystreamBytes = new byte[4];
            putU32(keystreamBytes, 0, keystream);

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

        // 清理上下文
        Arrays.fill(this.buf, (byte) 0);
        this.buflen = 0;
        Arrays.fill(this.state.LFSR, 0);
        this.state.R1 = 0;
        this.state.R2 = 0;
    }
}