#include "dbdcrypt.h"

#include <nerutils/log.h>

#include <algorithm>
#include <cstring>

#include <zlib.h>

#include <openssl/types.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/buffer.h>

std::string DBDCrypt::decrypt(const std::string& data, const std::string& accessKey, PayloadType* outType)
{
    if (outType) *outType = NONE;

    if (data.starts_with("DbdDAQEB"))
    {
        if (outType && *outType == NONE) *outType = TYPE_1;
        return decType1(data, accessKey, outType);
    }

    if (data.starts_with("DbdDAgAC"))
    {
        if (outType && *outType == NONE) *outType = TYPE_2;
        return decType2(data, accessKey, outType);
    }

    if (data.starts_with("DbdDAwAC"))
    {
        if (outType && *outType == NONE) *outType = TYPE_3;
        return decType3(data, accessKey, outType);
    }

    Log::warning("Attempted to decrypt non encrypted string");
    return data;
}

std::string DBDCrypt::encrypt(const std::string& data, const std::string& accessKey, PayloadType type,
                              std::string keyId)
{
    if (type == TYPE_1)
    {
        auto compressed = zLibCompress(data); // Type 1 is UTF-8
        if (compressed.empty()) return "";

        uint32_t rawSize = (uint32_t)data.size();
        std::vector<uint8_t> fullPayload(4);
        std::memcpy(fullPayload.data(), &rawSize, 4);

        fullPayload.insert(fullPayload.end(), compressed.begin(), compressed.end());
        return "DbdDAQEB" + b64Enc(fullPayload);
    }

    if (type == TYPE_2)
    {
        auto transformedKey = transformCDNKey(CDN_KEY_BASE64);

        std::string utf16Data = utf8ToUtf16(data);
        for (char& c : utf16Data)
            c = (char)((unsigned char)c - 1);

        std::vector<uint8_t> padded(utf16Data.begin(), utf16Data.end());

        int padLen = 16 - (padded.size() % 16);
        if (padLen < 16) padded.insert(padded.end(), padLen, 0);

        auto encryptedBody = aesECBEncrypt(padded, transformedKey);
        if (encryptedBody.empty()) return "";

        std::vector<uint8_t> fullPayload(encryptedBody.begin(), encryptedBody.end());
        return "DbdDAgAC" + b64Enc(fullPayload);
    }

    if (type == TYPE_3)
    {
        auto decodedKey = b64Dec(accessKey);
        if (decodedKey.empty()) return "";

        std::string utf16Data = utf8ToUtf16(data);
        for (char& c : utf16Data)
            c = (char)((unsigned char)c - 1);

        std::vector<uint8_t> padded(utf16Data.begin(), utf16Data.end());

        int padLen = 16 - (padded.size() % 16);
        if (padLen < 16) padded.insert(padded.end(), padLen, 0);

        auto encryptedBody = aesECBEncrypt(padded, decodedKey);

        if (encryptedBody.empty()) return "";

        std::string shiftedId = shiftKeyID(keyId, -1);
        std::vector<uint8_t> fullData(shiftedId.begin(), shiftedId.end());
        fullData.push_back(0); // Null terminator
        fullData.insert(fullData.end(), (uint8_t*)encryptedBody.data(),
                        (uint8_t*)encryptedBody.data() + encryptedBody.size());

        return "DbdDAwAC" + b64Enc(fullData);
    }

    return std::string();
}

std::string DBDCrypt::decType1(const std::string& data, const std::string& key, PayloadType* outType)
{
    if (data.length() < 8) return data;
    auto decoded = b64Dec(data.substr(8));
    if (decoded.size() < 4)
    {
        Log::error("Type 1 base64 too short ({})", decoded.size());
        return "";
    }

    std::vector<uint8_t> body(decoded.begin() + 4, decoded.end());
    std::string decompressed = zlibDecompress(body);

    if (decompressed.length() >= 2 && decompressed[1] == '\0') decompressed = utf16ToUtf8(decompressed);

    if (decompressed.starts_with("DbdD")) return decrypt(decompressed, key, outType);
    return decompressed;
}

std::string DBDCrypt::decType2(const std::string& data, const std::string& key, PayloadType* outType)
{
    if (data.length() < 8) return data;

    auto decoded = b64Dec(data.substr(8));
    if (decoded.empty()) return "";

    std::vector<uint8_t> body = decoded;

    auto transformedKey = transformCDNKey(CDN_KEY_BASE64);

    std::string decrypted = aesECBDecrypt(body, transformedKey);
    if (decrypted.empty()) return "";

    for (char& c : decrypted)
        c = (char)((unsigned char)c + 1);

    if (decrypted.length() >= 2 && decrypted[1] == '\0') decrypted = utf16ToUtf8(decrypted);

    decrypted.erase(std::remove(decrypted.begin(), decrypted.end(), (char)0x01), decrypted.end());
    decrypted.erase(std::remove(decrypted.begin(), decrypted.end(), (char)0x00), decrypted.end());

    for (size_t offset : {0ULL, 4ULL})
    {
        if (offset + 1 < decrypted.size() && (unsigned char)decrypted[offset] == 0x78)
        {
            std::vector<uint8_t> zlibPart((uint8_t*)decrypted.data() + offset,
                                          (uint8_t*)decrypted.data() + decrypted.size());
            std::string decompressed = zlibDecompress(zlibPart);
            if (!decompressed.empty()) return decompressed;
        }
    }

    if (decrypted.starts_with("DbdD")) return decrypt(decrypted, key, outType);
    return decrypted;
}

std::string DBDCrypt::decType3(const std::string& data, const std::string& key, PayloadType* outType)
{
    if (data.length() < 8) return data;

    auto rawKey = b64Dec(key);
    auto decoded = b64Dec(data.substr(8));
    if (decoded.empty()) return "";

    auto it = std::find(decoded.begin(), decoded.end(), 0);
    if (it == decoded.end()) return "";

    std::vector<uint8_t> body(it + 1, decoded.end());

    std::string decrypted = aesECBDecrypt(body, rawKey);
    if (decrypted.empty())
    {
        Log::error("AES decryption failed (body size: {})", body.size());
        return "";
    }

    for (char& c : decrypted)
        c = (char)((unsigned char)c + 1);

    if (decrypted.length() >= 2 && decrypted.at(1) == '\0') decrypted = utf16ToUtf8(decrypted);

    decrypted.erase(std::remove(decrypted.begin(), decrypted.end(), (char)0x01), decrypted.end());
    decrypted.erase(std::remove(decrypted.begin(), decrypted.end(), (char)0x00), decrypted.end());

    for (size_t offset : {0ULL, 4ULL})
    {
        if (offset + 1 < decrypted.size() && (unsigned char)decrypted[offset] == 0x78)
        {
            //Log::verbose("nested zlib at offset {}", offset);
            std::vector<uint8_t> zlibPart((uint8_t*)decrypted.data() + offset,
                                          (uint8_t*)decrypted.data() + decrypted.size());
            std::string decompressed = zlibDecompress(zlibPart);
            if (!decompressed.empty())
            {
                //Log::verbose("nested zlib decompressed, size: {}", decompressed.length());
                return decompressed;
            }
        }
    }

    if (decrypted.starts_with("DbdD")) return decrypt(decrypted, key, outType);
    return decrypted;
}

std::string DBDCrypt::aesECBDecrypt(const std::vector<uint8_t>& cipherText, const std::vector<uint8_t>& key)
{
    if (key.size() < 32) return "";

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    EVP_DecryptInit_ex(ctx, EVP_aes_256_ecb(), NULL, key.data(), NULL);
    EVP_CIPHER_CTX_set_padding(ctx, 0);

    std::string plaintext;
    plaintext.resize(cipherText.size());

    int len = 0;
    if (EVP_DecryptUpdate(ctx, (unsigned char*)plaintext.data(), &len, cipherText.data(), (int)cipherText.size()) != 1)
    {
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }
    int outLen = len;

    int finalLen = 0;
    if (EVP_DecryptFinal_ex(ctx, (unsigned char*)plaintext.data() + len, &finalLen) != 1)
    {
        // ignore?
    }
    outLen += finalLen;

    EVP_CIPHER_CTX_free(ctx);
    plaintext.resize(outLen);
    return plaintext;
}

std::string DBDCrypt::aesECBEncrypt(const std::vector<uint8_t>& plainText, const std::vector<uint8_t>& key)
{
    if (key.size() < 32) return "";

    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    EVP_EncryptInit_ex(ctx, EVP_aes_256_ecb(), NULL, key.data(), NULL);
    EVP_CIPHER_CTX_set_padding(ctx, 0);

    std::string ciphertext;
    ciphertext.resize(plainText.size() + 16);

    int len = 0;
    if (EVP_EncryptUpdate(ctx, (unsigned char*)ciphertext.data(), &len, plainText.data(), (int)plainText.size()) != 1)
    {
        EVP_CIPHER_CTX_free(ctx);
        return "";
    }
    int outLen = len;

    int finalLen = 0;
    if (EVP_EncryptFinal_ex(ctx, (unsigned char*)ciphertext.data() + outLen, &finalLen) == 1)
    {
        outLen += finalLen;
    }

    EVP_CIPHER_CTX_free(ctx);
    ciphertext.resize(outLen);
    return ciphertext;
}

std::vector<uint8_t> DBDCrypt::b64Dec(const std::string& input)
{
    std::string in = input;
    std::replace(in.begin(), in.end(), '-', '+');
    std::replace(in.begin(), in.end(), '_', '/');

    BIO *bio, *b64;
    int decodeLen = (int)in.length();
    std::vector<uint8_t> buffer(decodeLen);

    bio = BIO_new_mem_buf(in.data(), decodeLen);
    b64 = BIO_new(BIO_f_base64());
    bio = BIO_push(b64, bio);

    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
    int len = BIO_read(bio, buffer.data(), decodeLen);
    BIO_free_all(bio);

    if (len < 0) return {};
    buffer.resize(len);
    return buffer;
}

std::string DBDCrypt::b64Enc(const std::vector<uint8_t>& input)
{
    if (input.empty()) return "";

    BIO *bio, *b64;
    BUF_MEM* bufferPtr;

    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new(BIO_s_mem());
    bio = BIO_push(b64, bio);

    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
    BIO_write(bio, input.data(), (int)input.size());
    BIO_flush(bio);
    BIO_get_mem_ptr(bio, &bufferPtr);

    std::string result(bufferPtr->data, bufferPtr->length);
    BIO_free_all(bio);

    return result;
}

std::string DBDCrypt::zlibDecompress(const std::vector<uint8_t>& compressed)
{
    if (compressed.empty()) return "";

    z_stream strm;
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;
    strm.avail_in = (uInt)compressed.size();
    strm.next_in = (Bytef*)compressed.data();

    if (inflateInit(&strm) != Z_OK) return "";

    std::string result;
    char buffer[32768];

    do
    {
        strm.avail_out = sizeof(buffer);
        strm.next_out = (Bytef*)buffer;
        int ret = inflate(&strm, Z_NO_FLUSH);
        if (ret != Z_OK && ret != Z_STREAM_END)
        {
            inflateEnd(&strm);
            return "";
        }
        result.append(buffer, sizeof(buffer) - strm.avail_out);
    } while (strm.avail_out == 0);

    inflateEnd(&strm);
    return result;
}

std::vector<uint8_t> DBDCrypt::zLibCompress(const std::string& data)
{
    if (data.empty()) return {};

    z_stream strm;
    strm.zalloc = Z_NULL;
    strm.zfree = Z_NULL;
    strm.opaque = Z_NULL;

    if (deflateInit(&strm, Z_DEFAULT_COMPRESSION) != Z_OK) return {};

    strm.avail_in = (uInt)data.size();
    strm.next_in = (Bytef*)data.data();

    std::vector<uint8_t> result;
    uint8_t buffer[32768];

    do
    {
        strm.avail_out = sizeof(buffer);
        strm.next_out = (Bytef*)buffer;
        deflate(&strm, Z_FINISH);
        result.insert(result.end(), buffer, buffer + (sizeof(buffer) - strm.avail_out));
    } while (strm.avail_out == 0);

    deflateEnd(&strm);
    return result;
}

std::vector<uint8_t> DBDCrypt::transformCDNKey(const std::string& b64CDNKey)
{
    auto encryptedKey = b64Dec(b64CDNKey);
    std::vector<uint8_t> uuidKey(32, 0);
    std::memcpy(uuidKey.data(), CDN_UUID, std::min((size_t)32, strlen(CDN_UUID)));

    std::string decrypted = aesECBDecrypt(encryptedKey, uuidKey);
    std::vector<uint8_t> finalKey(decrypted.begin(), decrypted.end());
    if (finalKey.size() > 32)
        finalKey.resize(32);
    else if (finalKey.size() < 32)
        finalKey.resize(32, 0);
    return finalKey;
}

std::string DBDCrypt::shiftKeyID(const std::string& id, int shift)
{
    std::string res = id;
    for (char& c : res)
        c = (char)((unsigned char)c + shift);
    return res;
}

std::string DBDCrypt::utf16ToUtf8(const std::string& utf16)
{
    if (utf16.empty()) return "";
    std::string utf8;
    for (size_t i = 0; i < utf16.length(); i += 2)
    {
        uint16_t cp = *(uint16_t*)(utf16.data() + i);
        if (cp == 0) break;
        if (cp < 0x80)
            utf8 += (char)cp;
        else if (cp < 0x800)
        {
            utf8 += (char)(0xC0 | (cp >> 6));
            utf8 += (char)(0x80 | (cp & 0x3F));
        }
        else
        {
            utf8 += (char)(0xE0 | (cp >> 12));
            utf8 += (char)(0x80 | ((cp >> 6) & 0x3F));
            utf8 += (char)(0x80 | (cp & 0x3F));
        }
    }
    return utf8;
}

std::string DBDCrypt::utf8ToUtf16(const std::string& utf8)
{
    if (utf8.empty()) return "";
    std::string utf16;
    for (size_t i = 0; i < utf8.length();)
    {
        uint32_t cp = 0;
        unsigned char c = utf8[i];
        if (c < 0x80)
        {
            cp = c;
            i += 1;
        }
        else if (c < 0xE0)
        {
            cp = ((c & 0x1F) << 6) | (utf8[i + 1] & 0x3F);
            i += 2;
        }
        else if (c < 0xF0)
        {
            cp = ((c & 0x0F) << 12) | ((utf8[i + 1] & 0x3F) << 6) | (utf8[i + 2] & 0x3F);
            i += 3;
        }
        else
        {
            cp = ((c & 0x07) << 18) | ((utf8[i + 1] & 0x3F) << 12) | ((utf8[i + 2] & 0x3F) << 6) | (utf8[i + 3] & 0x3F);
            i += 4;
        }

        if (cp < 0x10000)
        {
            utf16.push_back((char)(cp & 0xFF));
            utf16.push_back((char)(cp >> 8));
        }
        else
        {
            cp -= 0x10000;
            uint16_t high = (uint16_t)(0xD800 | (cp >> 10));
            uint16_t low = (uint16_t)(0xDC00 | (cp & 0x3FF));

            utf16.push_back((char)(high & 0xFF));
            utf16.push_back((char)(high >> 8));
            utf16.push_back((char)(low & 0xFF));
            utf16.push_back((char)(low >> 8));
        }
    }
    return utf16;
}