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

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

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);
}

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;
}