#pragma once
#include "squirrel.h"
#include "sqstdaux.h"
#include "sqstdblob.h"
#include "sqstdio.h"
#include "sqstdmath.h"
#include "sqstdstring.h"
#include "sqstdsystem.h"

#include <iostream>
#include <fstream>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>


// 错误信息打印函数
void print_openssl_error() {
    char err_buf[256];
    ERR_error_string_n(ERR_get_error(), err_buf, sizeof(err_buf));
    std::cerr << "OpenSSL 错误: " << err_buf << std::endl;
}

// 生成 RSA 密钥对并保存到文件
// 生成 RSA 密钥对（公钥：X.509/PKCS#8 格式，私钥：PKCS#1 格式）
bool generate_rsa_key(const std::string& pub_key_path,
    const std::string& pri_key_path,
    int key_bits = 2048) {
    // 1. 创建 RSA 对象并生成密钥对
    RSA* rsa = RSA_new();
    BIGNUM* bn = BN_new();
    if (!rsa || !bn) {
        std::cerr << "创建 RSA/BIGNUM 失败" << std::endl;
        RSA_free(rsa);
        BN_free(bn);
        return false;
    }

    // 设置公钥指数（65537）
    if (BN_set_word(bn, RSA_F4) != 1) {
        std::cerr << "设置 BIGNUM 失败" << std::endl;
        print_openssl_error();
        RSA_free(rsa);
        BN_free(bn);
        return false;
    }

    // 生成 RSA 密钥对
    if (RSA_generate_key_ex(rsa, key_bits, bn, nullptr) != 1) {
        std::cerr << "生成 RSA 密钥对失败" << std::endl;
        print_openssl_error();
        RSA_free(rsa);
        BN_free(bn);
        return false;
    }

    // 2. 封装 RSA 到 EVP_PKEY（用于生成 X.509/PKCS#8 公钥）
    EVP_PKEY* evp_key = EVP_PKEY_new();
    if (!evp_key || EVP_PKEY_assign_RSA(evp_key, rsa) != 1) {
        std::cerr << "封装 EVP_PKEY 失败" << std::endl;
        print_openssl_error();
        EVP_PKEY_free(evp_key);
        RSA_free(rsa);
        BN_free(bn);
        return false;
    }

    // 3. 保存 X.509/PKCS#8 格式公钥（-----BEGIN PUBLIC KEY-----）
    FILE* pub_file = fopen(pub_key_path.c_str(), "w");
    if (!pub_file) {
        std::cerr << "打开公钥文件失败: " << pub_key_path << std::endl;
        EVP_PKEY_free(evp_key);
        BN_free(bn);
        return false;
    }
    if (PEM_write_PUBKEY(pub_file, evp_key) != 1) {
        std::cerr << "写入 X.509 公钥失败" << std::endl;
        print_openssl_error();
        fclose(pub_file);
        EVP_PKEY_free(evp_key);
        BN_free(bn);
        return false;
    }
    fclose(pub_file);

    // 4. 保存 PKCS#1 格式私钥（-----BEGIN RSA PRIVATE KEY-----）
    FILE* pri_file = fopen(pri_key_path.c_str(), "w");
    if (!pri_file) {
        std::cerr << "打开私钥文件失败: " << pri_key_path << std::endl;
        EVP_PKEY_free(evp_key);
        BN_free(bn);
        return false;
    }
    if (PEM_write_RSAPrivateKey(pri_file, rsa, nullptr, nullptr, 0, nullptr, nullptr) != 1) {
        std::cerr << "写入私钥失败" << std::endl;
        print_openssl_error();
        fclose(pri_file);
        EVP_PKEY_free(evp_key);
        BN_free(bn);
        return false;
    }
    fclose(pri_file);

    // 5. 释放资源（EVP_PKEY_free 会自动释放内部的 RSA 对象）
    EVP_PKEY_free(evp_key);
    BN_free(bn);

    std::cout << "RSA 密钥生成成功！" << std::endl;
    std::cout << "公钥（X.509/PKCS#8 格式）: " << pub_key_path << std::endl;
    std::cout << "私钥（PKCS#1 格式）: " << pri_key_path << std::endl;
    return true;
}


static SQInteger Rsa_GenerateKey(HSQUIRRELVM v)
{
    // 初始化 OpenSSL 错误处理
    ERR_load_crypto_strings();
    OpenSSL_add_all_algorithms();

    // 生成 2048 位 RSA 密钥对（推荐使用 2048/4096 位）
    bool ret = generate_rsa_key("public.pem", "private.pem", 2048);

    // 清理 OpenSSL 资源
    EVP_cleanup();
    ERR_free_strings();

    return 0;
}


// 从文件加载私钥（PKCS#1 格式）
RSA* load_rsa_private_key(const std::string& pri_key_path) {
    FILE* pri_file = fopen(pri_key_path.c_str(), "r");
    if (!pri_file) {
        std::cerr << "打开私钥文件失败: " << pri_key_path << std::endl;
        return nullptr;
    }

    // 加载 PKCS#1 格式私钥
    RSA* rsa = PEM_read_RSAPrivateKey(pri_file, nullptr, nullptr, nullptr);
    fclose(pri_file);

    if (!rsa) {
        std::cerr << "加载私钥失败" << std::endl;
        print_openssl_error();
    }
    return rsa;
}

// 新增：十六进制字符串转字节数组（核心补全）
bool hex_string_to_bytes(const std::string& hex_str, std::vector<unsigned char>& bytes) {
    bytes.clear();
    // 检查长度是否为偶数（十六进制字符串必须2字符=1字节）
    if (hex_str.length() % 2 != 0) {
        std::cerr << "十六进制字符串长度为奇数，无效" << std::endl;
        return false;
    }

    for (size_t i = 0; i < hex_str.length(); i += 2) {
        std::string byte_str = hex_str.substr(i, 2);
        // 十六进制转字节（处理0-9/a-f/A-F）
        unsigned char byte = static_cast<unsigned char>(strtol(byte_str.c_str(), nullptr, 16));
        bytes.push_back(byte);
    }
    return true;
}
bool rsa_private_encrypt(const std::vector<unsigned char>& plaintext_bytes, RSA* rsa_pri, std::string& ciphertext) {
    if (!rsa_pri || plaintext_bytes.empty()) return false;

    int rsa_len = RSA_size(rsa_pri);
    unsigned char* enc_buf = new unsigned char[rsa_len];

    // 加密字节数组（而非字符串的ASCII）
    int enc_len = RSA_private_encrypt(
        plaintext_bytes.size(),                // 明文长度（字节数）
        plaintext_bytes.data(),                // 明文字节数组
        enc_buf,                               // 加密结果缓冲区
        rsa_pri,                               // RSA私钥
        RSA_PKCS1_PADDING                      // 填充方式（和Java一致）
    );

    if (enc_len > 0) {
        ciphertext = std::string((char*)enc_buf, enc_len);
    } else {
        std::cerr << "私钥加密失败（2048位RSA最多加密245字节）" << std::endl;
        print_openssl_error();
    }

    delete[] enc_buf;
    return enc_len > 0;
}


static SQInteger Rsa_Private_Encrypt(HSQUIRRELVM v)
{
    const SQChar* hex_str;
    sq_getstring(v, 2, &hex_str); // 获取Squirrel传入的十六进制字符串

    // 1. 加载RSA私钥
    RSA* rsa_pri = load_rsa_private_key("privatekey.pem");
    if (!rsa_pri) {
        sq_pushnull(v);
        return 1;
    }

    // 2. 核心：十六进制字符串转字节数组（对齐Java的toByteArray）
    std::vector<unsigned char> plaintext_bytes;
    if (!hex_string_to_bytes(std::string(hex_str), plaintext_bytes)) {
        RSA_free(rsa_pri);
        sq_pushnull(v);
        return 1;
    }

    // 3. 加密字节数组
    std::string ciphertext;
    bool success = rsa_private_encrypt(plaintext_bytes, rsa_pri, ciphertext);

    // 4. 返回加密后的字节串（Squirrel中以字符串形式承载二进制数据）
    if (success) {
        sq_pushstring(v, ciphertext.c_str(), ciphertext.length());
    } else {
        sq_pushnull(v);
    }

    // 释放资源
    RSA_free(rsa_pri);
    return 1;
}

// ====================== 辅助函数 - 注册Squirrel函数 ======================
static SQInteger register_Rsa_func(HSQUIRRELVM v, SQFUNCTION f, const char* fname) {
    sq_pushroottable(v);
    sq_pushstring(v, fname, -1);
    sq_newclosure(v, f, 0);
    sq_newslot(v, -3, SQFalse);
    sq_pop(v, 1); // 弹出根表
    return 0;
}

// ====================== 注册所有Dio函数 ======================
static void RegisterRsa(HSQUIRRELVM v) {

    register_Rsa_func(v, Rsa_GenerateKey, _SC("Sq_Rsa_GenerateKey"));
    //使用私钥签名
    register_Rsa_func(v, Rsa_Private_Encrypt, _SC("Sq_Rsa_Private_Encrypt"));
}