#include "diality_secure_code.h"

#include <QtEndian>
#include <cstring>
#include <QDebug>

#include <openssl/evp.h>
#include <openssl/rand.h>

namespace diality {

static constexpr int kAesKeyBytes = 32;
static constexpr int kGcmIvBytes = 12;
static constexpr int kGcmTagBytes = 16;
static constexpr int kEdPkBytes = 32;
static constexpr int kEdSkBytes = 32;
static constexpr int kEdSigBytes = 64;

static QByteArray b64UrlNoPadEncodeQt(const QByteArray& bin)
{
#if (QT_VERSION >= QT_VERSION_CHECK(5, 15, 0))
    return bin.toBase64(QByteArray::Base64UrlEncoding | QByteArray::OmitTrailingEquals);
#else
    QByteArray b = bin.toBase64();
    b.replace('+', '-');
    b.replace('/', '_');
    while (!b.isEmpty() && b.endsWith('=')) b.chop(1);
    return b;
#endif
}

static bool b64UrlNoPadDecodeQt(const QByteArray& txt, QByteArray* binOut, QString* errOut)
{
    if (!binOut) {
        if (errOut) *errOut = "Internal error: binOut is null";
        return false;
    }
#if (QT_VERSION >= QT_VERSION_CHECK(5, 15, 0))
    QByteArray bin = QByteArray::fromBase64(txt, QByteArray::Base64UrlEncoding);
    if (bin.isEmpty() && !txt.isEmpty()) {
        if (errOut) *errOut = "Base64url decode failed";
        return false;
    }
    *binOut = bin;
    return true;
#else
    QByteArray t = txt;
    t.replace('-', '+');
    t.replace('_', '/');
    while (t.size() % 4) t.append('=');
    QByteArray bin = QByteArray::fromBase64(t);
    if (bin.isEmpty() && !txt.isEmpty()) {
        if (errOut) *errOut = "Base64url decode failed";
        return false;
    }
    *binOut = bin;
    return true;
#endif
}

bool cryptoInit()
{
#if OPENSSL_VERSION_NUMBER < 0x10100000L
    OpenSSL_add_all_algorithms();
#endif
    return true;
}

static bool aes256gcm_encrypt(
    const QByteArray& key32,
    const unsigned char iv[kGcmIvBytes],
    const QByteArray& plaintext,
    QByteArray* ciphertextPlusTag,
    QString* err
)
{
    if (key32.size() != kAesKeyBytes) { if (err) *err = "Bad AES key size"; return false; }
    if (!ciphertextPlusTag) { if (err) *err = "Internal error"; return false; }

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) { if (err) *err = "EVP_CIPHER_CTX_new failed"; return false; }

    bool ok = false;
    QByteArray ct(plaintext.size(), '\0');
    int outLen = 0;
    int total = 0;
    unsigned char tag[kGcmTagBytes];

    do {
        if (EVP_EncryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1) { if (err) *err = "EncryptInit (cipher) failed"; break; }
        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, kGcmIvBytes, nullptr) != 1) { if (err) *err = "Set IV len failed"; break; }
        if (EVP_EncryptInit_ex(ctx, nullptr, nullptr,
                               reinterpret_cast<const unsigned char*>(key32.constData()),
                               iv) != 1) { if (err) *err = "EncryptInit (key/iv) failed"; break; }

        if (EVP_EncryptUpdate(ctx,
                              reinterpret_cast<unsigned char*>(ct.data()),
                              &outLen,
                              reinterpret_cast<const unsigned char*>(plaintext.constData()),
                              plaintext.size()) != 1) { if (err) *err = "EncryptUpdate failed"; break; }
        total = outLen;

        if (EVP_EncryptFinal_ex(ctx,
                                reinterpret_cast<unsigned char*>(ct.data()) + total,
                                &outLen) != 1) { if (err) *err = "EncryptFinal failed"; break; }
        total += outLen;

        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, kGcmTagBytes, tag) != 1) { if (err) *err = "Get TAG failed"; break; }

        ct.truncate(total);
        QByteArray out;
        out.reserve(ct.size() + kGcmTagBytes);
        out.append(ct);
        out.append(reinterpret_cast<const char*>(tag), kGcmTagBytes);

        *ciphertextPlusTag = out;
        ok = true;
    } while (false);

    EVP_CIPHER_CTX_free(ctx);
    return ok;
}

static bool aes256gcm_decrypt(
    const QByteArray& key32,
    const unsigned char iv[kGcmIvBytes],
    const QByteArray& ciphertextPlusTag,
    QByteArray* plaintextOut,
    QString* err
)
{
    if (key32.size() != kAesKeyBytes) { if (err) *err = "Bad AES key size"; return false; }
    if (!plaintextOut) { if (err) *err = "Internal error"; return false; }
    if (ciphertextPlusTag.size() < kGcmTagBytes) { if (err) *err = "Ciphertext too short"; return false; }

    const int ctLen = ciphertextPlusTag.size() - kGcmTagBytes;
    const unsigned char* tag = reinterpret_cast<const unsigned char*>(ciphertextPlusTag.constData() + ctLen);

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) { if (err) *err = "EVP_CIPHER_CTX_new failed"; return false; }

    bool ok = false;
    QByteArray pt(ctLen, '\0');
    int outLen = 0;
    int total = 0;

    do {
        if (EVP_DecryptInit_ex(ctx, EVP_aes_256_gcm(), nullptr, nullptr, nullptr) != 1) { if (err) *err = "DecryptInit (cipher) failed"; break; }
        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_IVLEN, kGcmIvBytes, nullptr) != 1) { if (err) *err = "Set IV len failed"; break; }
        if (EVP_DecryptInit_ex(ctx, nullptr, nullptr,
                               reinterpret_cast<const unsigned char*>(key32.constData()),
                               iv) != 1) { if (err) *err = "DecryptInit (key/iv) failed"; break; }

        if (EVP_DecryptUpdate(ctx,
                              reinterpret_cast<unsigned char*>(pt.data()),
                              &outLen,
                              reinterpret_cast<const unsigned char*>(ciphertextPlusTag.constData()),
                              ctLen) != 1) { if (err) *err = "DecryptUpdate failed"; break; }
        total = outLen;

        if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_SET_TAG, kGcmTagBytes, const_cast<unsigned char*>(tag)) != 1) { if (err) *err = "Set TAG failed"; break; }

        const int fin = EVP_DecryptFinal_ex(ctx,
                                           reinterpret_cast<unsigned char*>(pt.data()) + total,
                                           &outLen);
        if (fin != 1) { if (err) *err = "Auth check failed (bad tag / tampered)"; break; }
        total += outLen;

        pt.truncate(total);
        *plaintextOut = pt;
        ok = true;
    } while (false);

    EVP_CIPHER_CTX_free(ctx);
    return ok;
}

static bool ed25519_sign_detached(
    const QByteArray& sk32,
    const QByteArray& msg,
    QByteArray* sig64,
    QString* err
)
{
    if (sk32.size() != kEdSkBytes) { if (err) *err = "Bad Ed25519 private key size"; return false; }
    if (!sig64) { if (err) *err = "Internal error"; return false; }

    EVP_PKEY* pkey = EVP_PKEY_new_raw_private_key(EVP_PKEY_ED25519, nullptr,
                                                  reinterpret_cast<const unsigned char*>(sk32.constData()),
                                                  kEdSkBytes);
    if (!pkey) { if (err) *err = "EVP_PKEY_new_raw_private_key failed"; return false; }

    EVP_MD_CTX* mctx = EVP_MD_CTX_new();
    if (!mctx) { EVP_PKEY_free(pkey); if (err) *err = "EVP_MD_CTX_new failed"; return false; }

    bool ok = false;
    size_t sigLen = kEdSigBytes;
    QByteArray sig(kEdSigBytes, '\0');

    do {
        if (EVP_DigestSignInit(mctx, nullptr, nullptr, nullptr, pkey) != 1) { if (err) *err = "DigestSignInit failed"; break; }
        if (EVP_DigestSign(mctx,
                           reinterpret_cast<unsigned char*>(sig.data()), &sigLen,
                           reinterpret_cast<const unsigned char*>(msg.constData()),
                           msg.size()) != 1) { if (err) *err = "DigestSign failed"; break; }
        if (sigLen != kEdSigBytes) { if (err) *err = "Unexpected signature size"; break; }
        *sig64 = sig;
        ok = true;
    } while (false);

    EVP_MD_CTX_free(mctx);
    EVP_PKEY_free(pkey);
    return ok;
}

static bool ed25519_verify_detached(
    const QByteArray& pk32,
    const QByteArray& msg,
    const QByteArray& sig64,
    QString* err
)
{
    if (pk32.size() != kEdPkBytes) { if (err) *err = "Bad Ed25519 public key size"; return false; }
    if (sig64.size() != kEdSigBytes) { if (err) *err = "Bad signature size"; return false; }

    EVP_PKEY* pkey = EVP_PKEY_new_raw_public_key(EVP_PKEY_ED25519, nullptr,
                                                 reinterpret_cast<const unsigned char*>(pk32.constData()),
                                                 kEdPkBytes);
    if (!pkey) { if (err) *err = "EVP_PKEY_new_raw_public_key failed"; return false; }

    EVP_MD_CTX* mctx = EVP_MD_CTX_new();
    if (!mctx) { EVP_PKEY_free(pkey); if (err) *err = "EVP_MD_CTX_new failed"; return false; }

    bool ok = false;
    do {
        if (EVP_DigestVerifyInit(mctx, nullptr, nullptr, nullptr, pkey) != 1) { if (err) *err = "DigestVerifyInit failed"; break; }
        const int rc = EVP_DigestVerify(mctx,
                                        reinterpret_cast<const unsigned char*>(sig64.constData()),
                                        sig64.size(),
                                        reinterpret_cast<const unsigned char*>(msg.constData()),
                                        msg.size());
        if (rc != 1) { if (err) *err = "Signature invalid"; break; }
        ok = true;
    } while (false);

    EVP_MD_CTX_free(mctx);
    EVP_PKEY_free(pkey);
    return ok;
}

QString buildSecureDatamatrixText(
    const QByteArray& plaintext,
    const QByteArray& encKey32,
    const QByteArray& signSk32,
    quint8 keyId
)
{
    if (encKey32.size() != kAesKeyBytes) return QString();
    if (signSk32.size() != kEdSkBytes) return QString();

    unsigned char iv[kGcmIvBytes];
    if (RAND_bytes(iv, sizeof(iv)) != 1) return QString();

    QByteArray cipherPlusTag;
    QString err;
    if (!aes256gcm_encrypt(encKey32, iv, plaintext, &cipherPlusTag, &err)) return QString();
    if (cipherPlusTag.size() > 65535) return QString();

    QByteArray container;
    container.reserve(4 + 1 + 1 + kGcmIvBytes + 2 + cipherPlusTag.size() + kEdSigBytes);

    container.append("DLY1", 4);
    container.append(char(0x01));     // Version
    container.append(char(keyId));    // KeyId
    container.append(reinterpret_cast<const char*>(iv), kGcmIvBytes);

    const quint16 cLen16 = quint16(cipherPlusTag.size());
    char lenBytes[2];
    qToLittleEndian<quint16>(cLen16, reinterpret_cast<uchar*>(lenBytes));
    container.append(lenBytes, 2);
    container.append(cipherPlusTag);

    QByteArray sig;
    if (!ed25519_sign_detached(signSk32, container, &sig, &err)) return QString();

    container.append(sig);

    const QByteArray b64 = b64UrlNoPadEncodeQt(container);
    return QString("DLY1.%1").arg(QString::fromLatin1(b64));
}

SecureCodeResult openSecureDatamatrixText(
    const QString& scannedText,
    const QByteArray& encKey32,
    const QByteArray& signPk32
)
{
    SecureCodeResult r;

    if (encKey32.size() != kAesKeyBytes) { r.error = "Bad encryption key size"; return r; }
    if (signPk32.size() != kEdPkBytes) { r.error = "Bad signing public key size"; return r; }

    const QString prefix = "DLY1.";
    if (!scannedText.startsWith(prefix)) { r.error = "Invalid prefix"; return r; }

    QByteArray bin;
    QString derr;
    if (!b64UrlNoPadDecodeQt(scannedText.mid(prefix.size()).toLatin1(), &bin, &derr)) { r.error = derr; return r; }

    const int minLen = 4 + 1 + 1 + kGcmIvBytes + 2 + kEdSigBytes;
    if (bin.size() < minLen) { r.error = "Container too short"; return r; }

    const int sigOff = bin.size() - kEdSigBytes;
    const QByteArray signedPart = bin.left(sigOff);
    const QByteArray sigPart = bin.mid(sigOff, kEdSigBytes);

    QString err;
    if (!ed25519_verify_detached(signPk32, signedPart, sigPart, &err)) {
        r.error = "Signature invalid (not Diality-generated)";
        return r;
    }

    const char* p = signedPart.constData();
    if (std::memcmp(p, "DLY1", 4) != 0) { r.error = "Bad magic"; return r; }

    const quint8 version = quint8(p[4]);
    if (version != 0x01) { r.error = "Unsupported version"; return r; }

    const int headerLen = 4 + 1 + 1 + kGcmIvBytes + 2;
    const unsigned char* iv = reinterpret_cast<const unsigned char*>(p + 6);
    const uchar* lenPtr = reinterpret_cast<const uchar*>(p + 6 + kGcmIvBytes);
    const quint16 cLen = qFromLittleEndian<quint16>(lenPtr);

    if (headerLen + int(cLen) != signedPart.size()) { r.error = "Length mismatch"; return r; }

    const QByteArray cipherPlusTag = signedPart.mid(headerLen, int(cLen));
    QByteArray plain;
    qDebug() << " ~ c ~ " << cipherPlusTag;
    qDebug() << " ~ p ~ " << plain;
    if (!aes256gcm_decrypt(encKey32, iv, cipherPlusTag, &plain, &err)) {
        r.error = "Decryption failed (tampered or wrong key)";
        return r;
    }

    r.ok = true;
    r.plaintext = plain;
    return r;
}

} // namespace diality