package com.zuc.zuc256;

import javacard.framework.JCSystem;
import javacard.framework.Util;

/**
 * 加密上下文类（Java Card兼容版本）
 */
public final class Zuc256EncryptCtx {
    Zuc256State state;
    byte[] buf;
    short buflen;

    // 构造函数 - 使用已分配的缓冲区
    public Zuc256EncryptCtx(Zuc256State state, byte[] buf) {
        this.state = state;
        this.buf = buf;
        this.buflen = 0;
    }

    // 构造函数 - 自动分配缓冲区
    public Zuc256EncryptCtx(Zuc256State state) {
        this.state = state;
        this.buf = JCSystem.makeTransientByteArray((short)4, JCSystem.CLEAR_ON_DESELECT);
        this.buflen = 0;
    }

    // 构造函数 - 完整初始化
    public Zuc256EncryptCtx() {
        this.state = new Zuc256State();
        this.buf = JCSystem.makeTransientByteArray((short)4, JCSystem.CLEAR_ON_DESELECT);
        this.buflen = 0;
    }

    // 初始化加密上下文
    public void init(byte[] key32, short keyOff, short keyLen, byte[] iv, short ivOff, short ivLen) {
        // 清空缓冲区
        Util.arrayFillNonAtomic(buf, (short)0, (short)buf.length, (byte)0);
        this.buflen = 0;
        // 初始化状态
        Zuc256Core.initState(this.state, key32, keyOff, keyLen, iv, ivOff, ivLen);
    }

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

        short currentInOff = inOff;
        short currentOutOff = outOff;

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

            // 复制数据到缓冲区
            Util.arrayCopyNonAtomic(in, currentInOff, this.buf, this.buflen, copy);
            this.buflen += copy;

            // 调整输入指针和长度
            currentInOff += copy;
            inlen -= copy;

            // 缓冲区已满，处理一个完整的4字节块
            if (this.buflen == 4) {
                short keystream = zuc256GenerateKeyword(this.state);
                short plain = getU16(this.buf, (short)0); // 改为16位操作
                putU16(out, currentOutOff, (short)(plain ^ keystream));

                // 重置缓冲区
                this.buflen = 0;
                Util.arrayFillNonAtomic(this.buf, (short)0, (short)this.buf.length, (byte)0);

                // 调整输出指针
                currentOutOff += 2; // 16位数据占2字节
            }
        }

        // 处理完整的2字节块（Java Card 16位操作）
        if (inlen > 0) {
            short fullBlocks = (short)(inlen / 2); // 16位块
            if (fullBlocks > 0) {
                short[] keystream = JCSystem.makeTransientShortArray(fullBlocks, JCSystem.CLEAR_ON_DESELECT);
                zuc256GenerateKeystream(this.state, fullBlocks, keystream);

                // 逐块异或加密
                for (short i = 0; i < fullBlocks; i++) {
                    short plain = getU16(in, (short)(currentInOff + (short)(i * 2)));
                    putU16(out, (short)(currentOutOff + (short)(i * 2)), (short)(plain ^ keystream[i]));
                }

                // 调整输入指针和长度
                short processed = (short)(fullBlocks * 2);
                currentInOff += processed;
                inlen -= processed;
            }

            // 缓存剩余不足2字节的数据
            if (inlen > 0) {
                Util.arrayCopyNonAtomic(in, currentInOff, this.buf, (short)0, inlen);
                this.buflen = inlen;
            }
        }
    }

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

        // 处理缓冲区中剩余的不足2字节数据
        if (this.buflen > 0) {
            short keystream = zuc256GenerateKeyword(this.state);
            byte[] keystreamBytes = JCSystem.makeTransientByteArray((short)2, JCSystem.CLEAR_ON_DESELECT);
            putU16(keystreamBytes, (short)0, keystream);

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

        // 清理上下文
        Util.arrayFillNonAtomic(this.buf, (short)0, (short)this.buf.length, (byte)0);
        this.buflen = 0;
        // 清理状态
        Util.arrayFillNonAtomic(this.state.LFSR, (short)0, (short)this.state.LFSR.length, (short)0);
        this.state.R1 = 0;
        this.state.R2 = 0;
    }

    // 16位数据读取（替代原32位实现）
    private short getU16(byte[] buf, short off) {
        return (short)(((buf[off] & 0xFF) << 8) | (buf[(short)(off + 1)] & 0xFF));
    }

    // 16位数据写入（替代原32位实现）
    private void putU16(byte[] buf, short off, short value) {
        buf[off] = (byte)((value >>> 8) & 0xFF);
        buf[(short)(off + 1)] = (byte)(value & 0xFF);
    }
}