#include "shortcrypto.h"
#include "tinyaes.h"
#include <QtDebug>
#include <QtMath>
#include <algorithm>

// Base85 字符集：85个可打印字符，避开容易在URL/JSON里出问题的字符
const char *ShortCrypto::BASE85_CHARS =
    "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!#$%&()*+-;<=>?@^_`{|}~";

// V3 字符集：纯大写字母（26字符），许可密钥只含 A-Z，28位
const char *ShortCrypto::V3_CHARS =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ";

namespace {
    QByteArray pkcs7Pad(const QByteArray &data, int blockSize) {
        int padLen = blockSize - (data.size() % blockSize);
        if (padLen == 0) padLen = blockSize;
        QByteArray result = data;
        result.append(QByteArray(padLen, static_cast<char>(padLen)));
        return result;
    }

    bool isPkcs7Valid(const QByteArray &data, int blockSize = 16) {
        if (data.isEmpty() || data.size() % blockSize != 0) return false;
        unsigned char padLen = static_cast<unsigned char>(data.at(data.size() - 1));
        if (padLen == 0 || padLen > static_cast<unsigned char>(blockSize)) return false;
        for (int i = 0; i < padLen; ++i)
            if (static_cast<unsigned char>(data.at(data.size() - 1 - i)) != padLen) return false;
        return true;
    }

    QByteArray pkcs7Unpad(const QByteArray &data) {
        if (data.isEmpty()) return QByteArray();
        unsigned char padLen = static_cast<unsigned char>(data.at(data.size() - 1));
        if (padLen == 0 || padLen > 16 || padLen > static_cast<unsigned char>(data.size())) return QByteArray();
        return data.left(data.size() - padLen);
    }
}

QString ShortCrypto::encrypt(const QByteArray &plaintext, const QByteArray &key, const QByteArray &iv)
{
    if (key.size() != 16) {
        qWarning() << "ShortCrypto::encrypt: key must be 16 bytes for AES-128";
        return QString();
    }

    QByteArray adjustedIv = iv;
    if (adjustedIv.isEmpty()) {
        adjustedIv = QByteArray(16, '\0');  // 默认全零IV（正式环境建议用随机IV并附带在密文前）
    } else if (adjustedIv.size() != 16) {
        adjustedIv = adjustedIv.left(16);
        adjustedIv.append(QByteArray(16 - adjustedIv.size(), '\0'));
    }

    // 1. PKCS7 填充
    QByteArray padded = pkcs7Pad(plaintext, 16);

    // 2. AES-128-CBC 加密（使用 Tiny-AES-C，纯软件实现，无 MSVC 兼容性问题）
    QByteArray encrypted = padded;
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, reinterpret_cast<const uint8_t*>(key.constData()),
                    reinterpret_cast<const uint8_t*>(adjustedIv.constData()));
    AES_CBC_encrypt_buffer(&ctx, reinterpret_cast<uint8_t*>(encrypted.data()), encrypted.size());

    // 3. Base85 编码，缩短字符数
    QByteArray encoded = base85Encode(encrypted);
    return QString::fromLatin1(encoded);
}

QByteArray ShortCrypto::decrypt(const QString &ciphertext, const QByteArray &key, const QByteArray &iv)
{
    if (key.size() != 16) {
        qWarning() << "ShortCrypto::decrypt: key must be 16 bytes for AES-128";
        return QByteArray();
    }

    QByteArray adjustedIv = iv;
    if (adjustedIv.isEmpty()) {
        adjustedIv = QByteArray(16, '\0');
    } else if (adjustedIv.size() != 16) {
        adjustedIv = adjustedIv.left(16);
        adjustedIv.append(QByteArray(16 - adjustedIv.size(), '\0'));
    }

    // 1. Base85 解码
    QByteArray encrypted = base85Decode(ciphertext);
    if (encrypted.isEmpty()) {
        qWarning() << "ShortCrypto::decrypt: base85 decode failed";
        return QByteArray();
    }

    if (encrypted.size() % 16 != 0) {
        qWarning() << "ShortCrypto::decrypt: decoded data size not aligned to block size:" << encrypted.size();
        return QByteArray();
    }

    // 2. AES-128-CBC 解密
    QByteArray decrypted = encrypted;
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, reinterpret_cast<const uint8_t*>(key.constData()),
                    reinterpret_cast<const uint8_t*>(adjustedIv.constData()));
    AES_CBC_decrypt_buffer(&ctx, reinterpret_cast<uint8_t*>(decrypted.data()), decrypted.size());

    // 3. PKCS7 验证和移除填充
    if (!isPkcs7Valid(decrypted)) {
        qWarning() << "ShortCrypto::decrypt: PKCS7 validation failed, wrong key or corrupted data";
        return QByteArray();
    }

    return pkcs7Unpad(decrypted);
}

QString ShortCrypto::encryptRaw(const QByteArray &plaintext, const QByteArray &key)
{
    if (key.size() != 16) {
        qWarning() << "ShortCrypto::encryptRaw: key must be 16 bytes for AES-128";
        return QString();
    }
    if (plaintext.isEmpty() || plaintext.size() % 16 != 0) {
        qWarning() << "ShortCrypto::encryptRaw: plaintext size must be a non-zero multiple of 16, got"
                   << plaintext.size();
        return QString();
    }

    // 单块时全零IV的CBC等价于ECB；多块时为CBC链接
    QByteArray adjustedIv(16, '\0');

    QByteArray encrypted = plaintext;
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, reinterpret_cast<const uint8_t*>(key.constData()),
                    reinterpret_cast<const uint8_t*>(adjustedIv.constData()));
    AES_CBC_encrypt_buffer(&ctx, reinterpret_cast<uint8_t*>(encrypted.data()), encrypted.size());

    // 纯大写字母编码：16字节密文 → 28字符（若允许数字换36字符集，长度降到25）
    QByteArray encoded = baseNEncode(encrypted, V3_CHARS, 26);
    return QString::fromLatin1(encoded);
}

QByteArray ShortCrypto::decryptRaw(const QString &ciphertext, const QByteArray &key)
{
    if (key.size() != 16) {
        qWarning() << "ShortCrypto::decryptRaw: key must be 16 bytes for AES-128";
        return QByteArray();
    }

    // 优先按纯字母表解码；若含表外字符（早期Base85版V3密钥）则回退Base85
    QByteArray encrypted = baseNDecode(ciphertext, V3_CHARS, 26);
    if (encrypted.isEmpty()) {
        encrypted = base85Decode(ciphertext);
    }
    if (encrypted.isEmpty()) {
        qWarning() << "ShortCrypto::decryptRaw: decode failed";
        return QByteArray();
    }
    if (encrypted.size() % 16 != 0) {
        qWarning() << "ShortCrypto::decryptRaw: decoded data size not aligned to block size:"
                   << encrypted.size();
        return QByteArray();
    }

    QByteArray adjustedIv(16, '\0');
    QByteArray decrypted = encrypted;
    struct AES_ctx ctx;
    AES_init_ctx_iv(&ctx, reinterpret_cast<const uint8_t*>(key.constData()),
                    reinterpret_cast<const uint8_t*>(adjustedIv.constData()));
    AES_CBC_decrypt_buffer(&ctx, reinterpret_cast<uint8_t*>(decrypted.data()), decrypted.size());
    return decrypted;
}

QByteArray ShortCrypto::baseNEncode(const QByteArray &data, const char *alphabet, int base)
{
    // 把整段大端字节序当作一个大整数，逐次除以base取余数
    QByteArray num = data;
    QByteArray remainders;

    bool isZero = false;
    while (!isZero) {
        isZero = true;
        unsigned int carry = 0;
        for (int i = 0; i < num.size(); ++i) {
            unsigned int cur = (carry << 8) | static_cast<unsigned char>(num.at(i));
            num[i] = static_cast<char>(cur / static_cast<unsigned int>(base));
            carry = cur % static_cast<unsigned int>(base);
            if (num.at(i) != 0) isZero = false;
        }
        if (isZero && remainders.isEmpty()) break; // 全零输入
        remainders.append(alphabet[static_cast<int>(carry)]);
        if (isZero) break;
    }
    if (remainders.isEmpty()) remainders.append(alphabet[0]);

    // 还原为高位在前，并用alphabet[0]补齐固定位数（前导零）
    QByteArray result = remainders;
    std::reverse(result.begin(), result.end());
    int minimalLen = static_cast<int>(qCeil(data.size() * 8.0 * qLn(2.0) / qLn(static_cast<double>(base))));
    if (result.size() < minimalLen) {
        result.prepend(QByteArray(minimalLen - result.size(), alphabet[0]));
    }
    return result;
}

QByteArray ShortCrypto::baseNDecode(const QString &text, const char *alphabet, int base)
{
    // 建立字符索引表
    int charToVal[256];
    memset(charToVal, -1, sizeof(charToVal));
    for (int i = 0; i < base; ++i) {
        charToVal[static_cast<unsigned char>(alphabet[i])] = i;
    }

    // Horner法则：num = num * base + v（大端字节序）
    QByteArray num(16, '\0');
    QByteArray input = text.toLatin1();
    for (int i = 0; i < input.size(); ++i) {
        int v = charToVal[static_cast<unsigned char>(input.at(i))];
        if (v < 0) {
            qWarning() << "ShortCrypto::baseNDecode: invalid character:" << QChar(input.at(i));
            return QByteArray();
        }
        unsigned int carry = static_cast<unsigned int>(v);
        for (int b = num.size() - 1; b >= 0; --b) {
            unsigned int cur = (static_cast<unsigned int>(static_cast<unsigned char>(num.at(b))) * static_cast<unsigned int>(base)) + carry;
            num[b] = static_cast<char>(cur & 0xFF);
            carry = cur >> 8;
        }
        if (carry != 0) {
            qWarning() << "ShortCrypto::baseNDecode: value overflows 128 bits";
            return QByteArray();
        }
    }

    // 保持16字节定宽（前导零有意义），调用方按固定块长处理
    return num;
}

QByteArray ShortCrypto::base85Encode(const QByteArray &data)
{
    QByteArray result;
    result.reserve((data.size() + 3) / 4 * 5);

    const unsigned char *ptr = reinterpret_cast<const unsigned char *>(data.constData());
    int len = data.size();

    while (len >= 4) {
        quint32 val = (quint32(ptr[0]) << 24) | (quint32(ptr[1]) << 16) |
                      (quint32(ptr[2]) << 8) | quint32(ptr[3]);
        char buf[5];
        for (int i = 4; i >= 0; --i) {
            buf[i] = BASE85_CHARS[val % 85];
            val /= 85;
        }
        result.append(buf, 5);
        ptr += 4;
        len -= 4;
    }

    // 处理剩余字节（1~3字节）
    if (len > 0) {
        quint32 val = 0;
        for (int i = 0; i < len; ++i) {
            val |= (quint32(ptr[i]) << (24 - i * 8));
        }
        char buf[5];
        for (int i = 4; i >= 0; --i) {
            buf[i] = BASE85_CHARS[val % 85];
            val /= 85;
        }
        result.append(buf, len + 1);  // n字节需要n+1个字符表示
    }

    return result;
}

QByteArray ShortCrypto::base85Decode(const QString &text)
{
    QByteArray input = text.toLatin1();
    QByteArray result;
    result.reserve(input.size() / 5 * 4);

    // 构建反向查找表
    int charToVal[256];
    memset(charToVal, -1, sizeof(charToVal));
    for (int i = 0; i < 85; ++i) {
        charToVal[static_cast<unsigned char>(BASE85_CHARS[i])] = i;
    }

    const unsigned char *ptr = reinterpret_cast<const unsigned char *>(input.constData());
    int len = input.size();
    int pos = 0;

    while (pos < len) {
        // 读取一个5字符块（或末尾的短块）
        int chunkLen = qMin(5, len - pos);
        quint32 val = 0;
        for (int i = 0; i < chunkLen; ++i) {
            int v = charToVal[ptr[pos + i]];
            if (v < 0) {
                qWarning() << "ShortCrypto::base85Decode: invalid character:" << QChar(ptr[pos + i]);
                return QByteArray();
            }
            val = val * 85 + v;
        }

        if (chunkLen == 5) {
            result.append(static_cast<char>((val >> 24) & 0xFF));
            result.append(static_cast<char>((val >> 16) & 0xFF));
            result.append(static_cast<char>((val >> 8) & 0xFF));
            result.append(static_cast<char>(val & 0xFF));
        } else {
            // 末尾短块：用 'u'(84) 填充到 5 个字符，再解码
            for (int i = chunkLen; i < 5; ++i)
                val = val * 85 + 84;

            char out[4];
            out[0] = static_cast<char>((val >> 24) & 0xFF);
            out[1] = static_cast<char>((val >> 16) & 0xFF);
            out[2] = static_cast<char>((val >> 8) & 0xFF);
            out[3] = static_cast<char>(val & 0xFF);

            int outBytes = chunkLen - 1;
            for (int i = 0; i < outBytes; ++i)
                result.append(out[i]);
        }
        pos += chunkLen;
    }

    return result;
}