#pragma once
#include <cstring>
#include <string>
#include <header/top_r3944realms_lib39_core_lang_ClassEncryptor.h>
#include <header/top_r3944realms_lib39_core_lang_EncryptedClassLoader.h>
#include "EnhancedEncryptionMagic.cpp"

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wsign-compare"
#ifdef _WIN32
#include <windows.h>
#include <wincrypt.h>
#include <cstdint>

#define JNIEXPORT __declspec(dllexport)
#else
#include <dlfcn.h>
    #define JNIEXPORT
#endif

/**
 * 简单的XOR加密/解密
 */
static  void xorEncrypt(jbyte* data, jsize dataLen, const char* key, int keyLen) {
    if (!data || !key || keyLen == 0) return;

    for (jsize i = 0; i < dataLen; i++) {
        data[i] ^= key[i % keyLen];
    }
}

static int safeStrlen(const char* str) {
    return str ? (int)strlen(str) : 0;
}

/**
 * 验证数据是否为有效的Java类文件
 * @param data 字节码
 * @param dataLen 字节码长度
 * @return 是否有效
 */
static bool isValidJavaClass(const jbyte* data, jsize dataLen) {
    // Java类文件魔数：0xCAFEBABE
    return dataLen >= 4 &&
           data[0] == (jbyte)0xCA &&
           data[1] == (jbyte)0xFE &&
           data[2] == (jbyte)0xBA &&
           data[3] == (jbyte)0xBE;
}

/**
 * 记录错误信息
 */
void logError(const char* message) {
#ifdef DEBUG
    std::cerr << "[JNI Error] " << message << std::endl;
#endif
}

/**
 * 安全的字符串转换
 */
std::string jstringToString(JNIEnv* env, jstring jstr) {
    if (!jstr) return "";

    const char* chars = env->GetStringUTFChars(jstr, nullptr);
    if (!chars) return "";

    std::string result(chars);
    env->ReleaseStringUTFChars(jstr, chars);
    return result;
}

uint32_t calculateChecksum(const jbyte* data, jsize length) {
    uint32_t checksum = 0;
    for (jsize i = 0; i < length; i++) {
        checksum += static_cast<uint8_t>(data[i]);
        checksum = (checksum << 1) | (checksum >> 31);  // 简单旋转
    }
    return checksum;
}


// ==================== JNI函数实现 ====================
using namespace EnhancedEncryptionMagic;

/*
 * Class:     top_r3944realms_lib39_core_lang_ClassEncryptor
 * Method:    decryptClass
 * Signature: ([BLjava/lang/String;)[B
 */
JNIEXPORT jbyteArray JNICALL Java_top_r3944realms_lib39_core_lang_ClassEncryptor_decryptClass
        (JNIEnv *env, jobject obj, jbyteArray encryptedData, jstring key) {

    if (!encryptedData || !key) {
        return nullptr;
    }

    jsize fileSize = env->GetArrayLength(encryptedData);
    if (fileSize == 0) {
        return nullptr;
    }
    JByteArrayGuard fileDataGuard(env, encryptedData);
    jbyte* fileData = fileDataGuard.get();
    if (!fileDataGuard.isValid()) {
        return nullptr;
    }
    std::string keyStr = jstringToString(env, key);
    size_t keyLen = keyStr.length();
    bool isValid = keyLen > 0;
    // 尝试从加密文件中提取数据
    jbyteArray result = EnhancedEncryptionMagic::extractFromEncryptedFile(
            env, fileData, fileSize, keyStr.c_str(), keyLen, xorEncrypt, &isValid);
    if (!isValid || !result) {
        // 如果不是有效的加密文件，返回原始数据
        return encryptedData;
    }

    // 验证解密后的数据是否为有效的Java类文件
    JByteArrayGuard resultGuard(env, result);
    jsize resultLen = resultGuard.size();
    jbyte* resultData = resultGuard.get();

    if (!resultGuard.isValid()) {
        env->DeleteLocalRef(result);
        return nullptr;
    }

    bool validClass = isValidJavaClass(resultData, resultLen);

    if (!validClass) {
        env->DeleteLocalRef(result);
        return nullptr;  // 解密后的数据不是有效的Java类
    }

    return result;
}


/*
 * Class:     top_r3944realms_lib39_core_lang_ClassEncryptor
 * Method:    encryptClass
 * Signature: ([BLjava/lang/String;)[B
 */
JNIEXPORT jbyteArray JNICALL Java_top_r3944realms_lib39_core_lang_ClassEncryptor_encryptClass
        (JNIEnv *env, jobject obj, jbyteArray classData, jstring key) {

    if (!classData || !key) {
        return nullptr;
    }

    JByteArrayGuard jGuard(env, classData);
    if (jGuard.isValid()) {
        return nullptr;
    }

    jsize dataLen = jGuard.size();
    jbyte* data = jGuard.get();

    const char* keyStr = env->GetStringUTFChars(key, nullptr);
    if (!keyStr) {
        return nullptr;
    }

    // 验证输入数据是否为有效的Java类文件
    if (!isValidJavaClass(data, dataLen)) {
        env->ReleaseStringUTFChars(key, keyStr);
        return nullptr;
    }

    int keyLen = safeStrlen(keyStr);

    // 使用增强版创建加密文件
    jbyteArray result = EnhancedEncryptionMagic::createEncryptedFile(
            env, data, dataLen, keyStr, keyLen, xorEncrypt);

    env->ReleaseStringUTFChars(key, keyStr);

    return result;
}

/*
 * Class:     top_r3944realms_lib39_core_lang_ClassEncryptor
 * Method:    isEncryptedFile
 * Signature: ([B)Z
 */
JNIEXPORT jboolean JNICALL Java_top_r3944realms_lib39_core_lang_ClassEncryptor_isEncryptedFile
        (JNIEnv *env, jobject obj, jbyteArray fileData) {

    if (!fileData) {
        return JNI_FALSE;
    }
    JByteArrayGuard dataGuard(env, fileData);
    if (!dataGuard.isValid()) {
        return JNI_FALSE;
    }

    if (dataGuard.size() < EnhancedEncryptionMagic::HEADER_SIZE) {
        return JNI_FALSE;
    }

    // 验证是否为加密文件
    bool isEncrypted = EnhancedEncryptionMagic::isEncryptedFile(dataGuard.get(), dataGuard.size());

    return isEncrypted ? JNI_TRUE : JNI_FALSE;
}

/*
 * Class:     top_r3944realms_lib39_core_lang_EncryptedClassLoader
 * Method:    decryptClass
 * Signature: ([BLjava/lang/String;)[B
 */
JNIEXPORT jbyteArray JNICALL Java_top_r3944realms_lib39_core_lang_EncryptedClassLoader_decryptClass
        (JNIEnv *env, jobject obj, jbyteArray encryptedData, jstring key) {
    return Java_top_r3944realms_lib39_core_lang_ClassEncryptor_decryptClass(env, obj, encryptedData, key);
}

// JNI库初始化和卸载函数
JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void* reserved) {
#ifdef _WIN32
    // Windows下需要初始化Winsock
    WSADATA wsaData;
    WSAStartup(MAKEWORD(2, 2), &wsaData);
#endif

    return JNI_VERSION_1_8;
}

JNIEXPORT void JNICALL JNI_OnUnload(JavaVM* vm, void* reserved) {
#ifdef _WIN32
    WSACleanup();
#endif
}
#pragma clang diagnostic pop