/*
 * Copyright (c) 2025 SimcuTeam. All rights reserved.
 * 遵循 MIT 许可证。
 */

package simcu::simapi.helpers

import std.collection.*
import std.random.*
import stdx.crypto.digest.*
import stdx.encoding.base64.*

/**
 * AES-256-CBC + PKCS7 加解密工具。
 *
 * 约定：
 *   - 密钥：SHA256(key 字符串) → 32 字节
 *   - 模式：AES-256-CBC，PKCS7 填充
 *   - IV：每次加密随机生成 16 字节，前置在密文前
 *   - 输出：Base64(IV(16) + 密文)
 *
 * 仓颉生态（stdx / soulsoft）均无现成 AES 实现，此处纯仓颉实现 FIPS-197 AES-256。
 */
public class SimApiAesUtil {
    private init() {}

    /**
     * AES 加密：Base64(随机IV + 密文)。
     * @param plainText 明文。
     * @param key 字符串密钥（SHA256 处理后为 256 位）。
     * @return Base64(IV + 密文)。
     */
    public static func encrypt(plainText: String, key: String): String {
        if (plainText.isEmpty()) {
            throw Exception("plainText 不能为空")
        }
        if (key.isEmpty()) {
            throw Exception("key 不能为空")
        }
        let keyBytes = processKey(key)
        let iv = generateIv()
        let padded = pkcs7Pad(plainText.toArray())
        let cipher = cbcEncrypt(padded, keyBytes, iv)

        // IV + 密文 → Base64
        var out = ArrayList<Byte>()
        for (b in iv) {
            out.add(toB(b))
        }
        for (b in cipher) {
            out.add(toB(b))
        }
        toBase64String(out.toArray())
    }

    /**
     * AES 解密。
     * @param cipherText Base64(IV + 密文)。
     * @param key 字符串密钥（与加密时相同）。
     * @return 明文。
     */
    public static func decrypt(cipherText: String, key: String): String {
        if (cipherText.isEmpty()) {
            throw Exception("cipherText 不能为空")
        }
        if (key.isEmpty()) {
            throw Exception("key 不能为空")
        }
        let all = fromBase64String(cipherText).getOrThrow { Exception("Base64 解码失败") }
        if (all.size < 32) {
            throw Exception("密文长度非法")
        }
        var iv = Array<UInt8>(16, repeat: 0)
        var cipher = Array<UInt8>(all.size - 16, repeat: 0)
        for (i in 0..16) {
            iv[i] = toU8(all[i])
        }
        for (i in 0..cipher.size) {
            cipher[i] = toU8(all[i + 16])
        }
        let keyBytes = processKey(key)
        let padded = cbcDecrypt(cipher, keyBytes, iv)
        let plain = pkcs7Unpad(padded)
        String.fromUtf8(toBytes(plain))
    }

    // ===== 密钥与 IV =====

    /// SHA256(key) → 32 字节密钥
    @OverflowWrapping
    private static func processKey(key: String): Array<UInt8> {
        let md = SHA256()
        md.write(key.toArray())
        let digest = md.finish()
        var out = Array<UInt8>(digest.size, repeat: 0)
        for (i in 0..digest.size) {
            out[i] = toU8(digest[i])
        }
        out
    }

    /// 16 字节随机 IV
    private static func generateIv(): Array<UInt8> {
        let rnd = Random()
        var iv = Array<UInt8>(16, repeat: 0)
        for (i in 0..16) {
            iv[i] = UInt8(rnd.nextUInt64() & 0xFFu64)
        }
        iv
    }

    // ===== PKCS7 填充 =====

    @OverflowWrapping
    private static func pkcs7Pad(data: Array<Byte>): Array<UInt8> {
        let padLen = 16 - (data.size % 16)
        var out = ArrayList<UInt8>()
        for (b in data) {
            out.add(toU8(b))
        }
        for (i in 0..padLen) {
            out.add(UInt8(padLen))
        }
        out.toArray()
    }

    @OverflowWrapping
    private static func pkcs7Unpad(data: Array<UInt8>): Array<UInt8> {
        if (data.size == 0) {
            return Array<UInt8>(0, repeat: 0)
        }
        let padLen = Int64(data[data.size - 1])
        if (padLen < 1 || padLen > 16) {
            throw Exception("PKCS7 填充非法")
        }
        data[0..data.size - padLen]
    }

    // ===== CBC 模式 =====

    @OverflowWrapping
    private static func cbcEncrypt(padded: Array<UInt8>, key: Array<UInt8>, iv: Array<UInt8>): Array<UInt8> {
        let roundKeys = keyExpansion(key)
        var prev = Array<UInt8>(16, repeat: 0)
        for (i in 0..16) {
            prev[i] = iv[i]
        }
        var out = ArrayList<UInt8>()
        let n = padded.size / 16
        for (block in 0..n) {
            var state = Array<UInt8>(16, repeat: 0)
            for (i in 0..16) {
                state[i] = padded[block * 16 + i] ^ prev[i]
            }
            let enc = aesEncryptBlock(state, roundKeys)
            for (b in enc) {
                out.add(b)
            }
            prev = enc
        }
        out.toArray()
    }

    @OverflowWrapping
    private static func cbcDecrypt(cipher: Array<UInt8>, key: Array<UInt8>, iv: Array<UInt8>): Array<UInt8> {
        let roundKeys = keyExpansion(key)
        var prev = Array<UInt8>(16, repeat: 0)
        for (i in 0..16) {
            prev[i] = iv[i]
        }
        var out = ArrayList<UInt8>()
        let n = cipher.size / 16
        for (block in 0..n) {
            var blockBytes = Array<UInt8>(16, repeat: 0)
            for (i in 0..16) {
                blockBytes[i] = cipher[block * 16 + i]
            }
            let dec = aesDecryptBlock(blockBytes, roundKeys)
            for (i in 0..16) {
                out.add(dec[i] ^ prev[i])
            }
            prev = blockBytes
        }
        out.toArray()
    }

    // ===== AES-256 核心（FIPS-197） =====

    /// S-box / 逆 S-box（运行时生成，避免手写 256 字节表出错）
    private static let sbox: Array<UInt8> = generateSbox()
    private static let invSbox: Array<UInt8> = generateInvSbox()

    @OverflowWrapping
    private static func generateSbox(): Array<UInt8> {
        var s = Array<UInt8>(256, repeat: 0)
        var p = 1u8
        var q = 1u8
        while (true) {
            // p *= 3（GF(2^8) 生成元遍历）
            p = p ^ (p << 1u8) ^ (if ((p & 0x80u8) != 0u8) { 0x1Bu8 } else { 0u8 })
            // q /= 3（等价乘以 0xF6）
            q = q ^ (q << 1u8)
            q = q ^ (q << 2u8)
            q = q ^ (q << 4u8)
            q = q ^ (if ((q & 0x80u8) != 0u8) { 0x09u8 } else { 0u8 })
            // 仿射变换
            let x = q ^ rotl8(q, 1) ^ rotl8(q, 2) ^ rotl8(q, 3) ^ rotl8(q, 4)
            s[Int64(p)] = x ^ 0x63u8
            if (p == 1u8) {
                break
            }
        }
        // p 序列遍历非零元素，s[0] 从未赋值：AES 规定 S(0) = 0x63（0 的逆为 0，仿射变换结果）
        s[0] = 0x63u8
        s
    }

    @OverflowWrapping
    private static func generateInvSbox(): Array<UInt8> {
        var inv = Array<UInt8>(256, repeat: 0)
        for (i in 0..256) {
            inv[Int64(sbox[i])] = UInt8(i)
        }
        inv
    }

    /// 8 位循环左移
    private static func rotl8(v: UInt8, n: Int64): UInt8 {
        UInt8(((UInt16(v) << n) | (UInt16(v) >> (8 - n))) & 0xFFu16)
    }

    /// GF(2^8) 乘以 2（xtime）
    @OverflowWrapping
    private static func xtime(a: UInt8): UInt8 {
        if ((a & 0x80u8) != 0u8) {
            (a << 1u8) ^ 0x1Bu8
        } else {
            a << 1u8
        }
    }

    /// GF(2^8) 通用乘法
    @OverflowWrapping
    private static func gfMul(a: UInt8, b: UInt8): UInt8 {
        var result = 0u8
        var aa = a
        var bb = b
        for (i in 0..8) {
            if ((bb & 1u8) != 0u8) {
                result = result ^ aa
            }
            let hi = (aa & 0x80u8) != 0u8
            aa = aa << 1u8
            if (hi) {
                aa = aa ^ 0x1Bu8
            }
            bb = bb >> 1u8
        }
        result
    }

    /// 密钥扩展：32 字节密钥 → 240 字节轮密钥（60 words，AES-256 共 15 轮）
    @OverflowWrapping
    private static func keyExpansion(key: Array<UInt8>): Array<UInt8> {
        var w = Array<UInt8>(240, repeat: 0)
        for (i in 0..32) {
            w[i] = key[i]
        }
        var rcon = 1u8
        for (i in 8..60) {
            var temp = Array<UInt8>(4, repeat: 0)
            for (j in 0..4) {
                temp[j] = w[(i - 1) * 4 + j]
            }
            if (i % 8 == 0) {
                // RotWord
                let t0 = temp[0]
                temp[0] = temp[1]
                temp[1] = temp[2]
                temp[2] = temp[3]
                temp[3] = t0
                // SubWord
                for (j in 0..4) {
                    temp[j] = sbox[Int64(temp[j])]
                }
                temp[0] = temp[0] ^ rcon
                rcon = xtime(rcon)
            } else if (i % 8 == 4) {
                for (j in 0..4) {
                    temp[j] = sbox[Int64(temp[j])]
                }
            }
            for (j in 0..4) {
                w[i * 4 + j] = w[(i - 8) * 4 + j] ^ temp[j]
            }
        }
        w
    }

    @OverflowWrapping
    private static func aesEncryptBlock(input: Array<UInt8>, roundKeys: Array<UInt8>): Array<UInt8> {
        var state = Array<UInt8>(16, repeat: 0)
        for (i in 0..16) {
            state[i] = input[i]
        }
        addRoundKey(state, roundKeys, 0)
        for (round in 1..14) {
            subBytes(state)
            shiftRows(state)
            mixColumns(state)
            addRoundKey(state, roundKeys, round)
        }
        subBytes(state)
        shiftRows(state)
        addRoundKey(state, roundKeys, 14)
        state
    }

    @OverflowWrapping
    private static func aesDecryptBlock(input: Array<UInt8>, roundKeys: Array<UInt8>): Array<UInt8> {
        var state = Array<UInt8>(16, repeat: 0)
        for (i in 0..16) {
            state[i] = input[i]
        }
        addRoundKey(state, roundKeys, 14)
        // 轮 13..1（逆序）
        for (i in 0..13) {
            let round = 13 - i
            invShiftRows(state)
            invSubBytes(state)
            addRoundKey(state, roundKeys, round)
            invMixColumns(state)
        }
        invShiftRows(state)
        invSubBytes(state)
        addRoundKey(state, roundKeys, 0)
        state
    }

    @OverflowWrapping
    private static func addRoundKey(state: Array<UInt8>, roundKeys: Array<UInt8>, round: Int64): Unit {
        for (i in 0..16) {
            state[i] = state[i] ^ roundKeys[round * 16 + i]
        }
    }

    @OverflowWrapping
    private static func subBytes(state: Array<UInt8>): Unit {
        for (i in 0..16) {
            state[i] = sbox[Int64(state[i])]
        }
    }

    @OverflowWrapping
    private static func invSubBytes(state: Array<UInt8>): Unit {
        for (i in 0..16) {
            state[i] = invSbox[Int64(state[i])]
        }
    }

    /// ShiftRows：行 r 循环左移 r 字节（列主序 state[i] = s[r + 4*c]）
    @OverflowWrapping
    private static func shiftRows(state: Array<UInt8>): Unit {
        var tmp = Array<UInt8>(16, repeat: 0)
        for (r in 0..4) {
            for (c in 0..4) {
                tmp[r + 4 * c] = state[r + 4 * ((c + r) % 4)]
            }
        }
        for (i in 0..16) {
            state[i] = tmp[i]
        }
    }

    /// InvShiftRows：行 r 循环右移 r 字节
    @OverflowWrapping
    private static func invShiftRows(state: Array<UInt8>): Unit {
        var tmp = Array<UInt8>(16, repeat: 0)
        for (r in 0..4) {
            for (c in 0..4) {
                tmp[r + 4 * c] = state[r + 4 * ((c - r % 4 + 4) % 4)]
            }
        }
        for (i in 0..16) {
            state[i] = tmp[i]
        }
    }

    /// MixColumns：每列乘固定矩阵 [[2,3,1,1],[1,2,3,1],[1,1,2,3],[3,1,1,2]]
    @OverflowWrapping
    private static func mixColumns(state: Array<UInt8>): Unit {
        for (c in 0..4) {
            let a0 = state[0 + 4 * c]
            let a1 = state[1 + 4 * c]
            let a2 = state[2 + 4 * c]
            let a3 = state[3 + 4 * c]
            state[0 + 4 * c] = gfMul(a0, 2u8) ^ gfMul(a1, 3u8) ^ a2 ^ a3
            state[1 + 4 * c] = a0 ^ gfMul(a1, 2u8) ^ gfMul(a2, 3u8) ^ a3
            state[2 + 4 * c] = a0 ^ a1 ^ gfMul(a2, 2u8) ^ gfMul(a3, 3u8)
            state[3 + 4 * c] = gfMul(a0, 3u8) ^ a1 ^ a2 ^ gfMul(a3, 2u8)
        }
    }

    /// InvMixColumns：每列乘逆矩阵 [[14,11,13,9],[9,14,11,13],[13,9,14,11],[11,13,9,14]]
    @OverflowWrapping
    private static func invMixColumns(state: Array<UInt8>): Unit {
        for (c in 0..4) {
            let a0 = state[0 + 4 * c]
            let a1 = state[1 + 4 * c]
            let a2 = state[2 + 4 * c]
            let a3 = state[3 + 4 * c]
            state[0 + 4 * c] = gfMul(a0, 14u8) ^ gfMul(a1, 11u8) ^ gfMul(a2, 13u8) ^ gfMul(a3, 9u8)
            state[1 + 4 * c] = gfMul(a0, 9u8) ^ gfMul(a1, 14u8) ^ gfMul(a2, 11u8) ^ gfMul(a3, 13u8)
            state[2 + 4 * c] = gfMul(a0, 13u8) ^ gfMul(a1, 9u8) ^ gfMul(a2, 14u8) ^ gfMul(a3, 11u8)
            state[3 + 4 * c] = gfMul(a0, 11u8) ^ gfMul(a1, 13u8) ^ gfMul(a2, 9u8) ^ gfMul(a3, 14u8)
        }
    }

    // ===== 字节转换 =====

    /// Byte → UInt8：负 Byte（>127 字节的补码）按位重解释为 0-255
    private static func toU8(b: Byte): UInt8 {
        if (b < 0) {
            UInt8(Int64(b) + 256)
        } else {
            UInt8(Int64(b))
        }
    }

    private static func toB(u: UInt8): Byte {
        let b: Byte = u
        b
    }

    private static func toBytes(data: Array<UInt8>): Array<Byte> {
        var out = Array<Byte>(data.size, repeat: 0)
        for (i in 0..data.size) {
            out[i] = toB(data[i])
        }
        out
    }
}
