#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/rsa.h>
#include <openssl/sha.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>

/**
 * @brief 计算文件的SHA256哈希值
 * @param filePath 文件路径
 * @param hash 输出哈希值（32字节）
 * @return 0表示成功，-1表示失败
 */
int calculateFileHash(const char *filePath, unsigned char *hash)
{
	if (!filePath || !hash) {
		return -1;
	}

	std::ifstream file(filePath, std::ios::binary);
	if (!file.is_open()) {
		std::cerr << "错误：无法打开文件 " << filePath << std::endl;
		return -1;
	}

	SHA256_CTX sha256;
	SHA256_Init(&sha256);

	char buffer[4096];
	while (file) {
		file.read(buffer, sizeof(buffer));
		std::streamsize bytesRead = file.gcount();
		if (bytesRead > 0) {
			SHA256_Update(&sha256, buffer, bytesRead);
		}
	}

	SHA256_Final(hash, &sha256);
	file.close();

	return 0;
}

/**
 * @brief 为指定文件创建数字签名
 * @param filePath 要签名的文件路径
 * @param privateKeyPath 私钥文件路径
 * @param signaturePath 签名文件保存路径
 * @return 0表示成功，-1表示失败
 */
int signFile(const char *filePath, const char *privateKeyPath, const char *signaturePath)
{
	if (!filePath || !privateKeyPath || !signaturePath) {
		std::cerr << "错误：参数不能为空" << std::endl;
		return -1;
	}

	// 加载私钥
	FILE *privateKeyFile = fopen(privateKeyPath, "r");
	if (!privateKeyFile) {
		std::cerr << "错误：无法打开私钥文件" << std::endl;
		return -1;
	}

	EVP_PKEY *privateKey = PEM_read_PrivateKey(privateKeyFile, nullptr, nullptr, nullptr);
	fclose(privateKeyFile);

	if (!privateKey) {
		std::cerr << "错误：无法读取私钥" << std::endl;
		return -1;
	}

	// 计算文件哈希
	unsigned char hash[SHA256_DIGEST_LENGTH];
	if (calculateFileHash(filePath, hash) != 0) {
		EVP_PKEY_free(privateKey);
		return -1;
	}

	// 创建签名上下文
	EVP_MD_CTX *ctx = EVP_MD_CTX_new();
	if (!ctx) {
		std::cerr << "错误：无法创建签名上下文" << std::endl;
		EVP_PKEY_free(privateKey);
		return -1;
	}

	// 初始化签名
	if (EVP_DigestSignInit(ctx, nullptr, EVP_sha256(), nullptr, privateKey) <= 0) {
		std::cerr << "错误：无法初始化签名" << std::endl;
		EVP_MD_CTX_free(ctx);
		EVP_PKEY_free(privateKey);
		return -1;
	}

	// 计算签名长度
	size_t signatureLen;
	if (EVP_DigestSign(ctx, nullptr, &signatureLen, hash, SHA256_DIGEST_LENGTH) <= 0) {
		std::cerr << "错误：无法计算签名长度" << std::endl;
		EVP_MD_CTX_free(ctx);
		EVP_PKEY_free(privateKey);
		return -1;
	}

	// 分配签名缓冲区
	std::vector<unsigned char> signature(signatureLen);

	// 执行签名
	if (EVP_DigestSign(ctx, signature.data(), &signatureLen, hash, SHA256_DIGEST_LENGTH) <= 0) {
		std::cerr << "错误：签名失败" << std::endl;
		EVP_MD_CTX_free(ctx);
		EVP_PKEY_free(privateKey);
		return -1;
	}

	// 保存签名到文件
	std::ofstream signatureFile(signaturePath, std::ios::binary);
	if (!signatureFile.is_open()) {
		std::cerr << "错误：无法创建签名文件" << std::endl;
		EVP_MD_CTX_free(ctx);
		EVP_PKEY_free(privateKey);
		return -1;
	}

	signatureFile.write(reinterpret_cast<const char *>(signature.data()), signatureLen);
	signatureFile.close();

	// 清理资源
	EVP_MD_CTX_free(ctx);
	EVP_PKEY_free(privateKey);

	std::cout << "文件签名成功！" << std::endl;
	std::cout << "签名文件保存到: " << signaturePath << std::endl;

	return 0;
}

#include <filesystem>
#include <vector>
#include <string>

/**
 * @brief 获取目录下所有文件
 * @param directory 目录路径
 * @return 文件路径列表
 */
std::vector<std::string> getFilesInDirectory(const std::string &directory)
{
	std::vector<std::string> files;

	try {
		for (const auto &entry : std::filesystem::directory_iterator(directory)) {
			if (entry.is_regular_file()) {
				// 跳过密钥文件和签名文件
				std::string filename = entry.path().filename().string();
				if (filename != "private_key.pem" && filename != "public_key.pem"
					&& filename.find(".sig") == std::string::npos) {
					files.push_back(entry.path().string());
				}
			}
		}
	} catch (const std::filesystem::filesystem_error &e) {
		std::cerr << "错误：无法读取目录 " << directory << ": " << e.what() << std::endl;
	}

	return files;
}

/**
 * @brief 显示使用说明
 */
void showUsage(const char *programName)
{
	std::cout << "用法: " << programName << " [文件路径]" << std::endl;
	std::cout << "参数说明:" << std::endl;
	std::cout << "  文件路径: 要签名的文件路径（可选）" << std::endl;
	std::cout << "  如果不提供参数，将签名当前目录下的所有文件" << std::endl;
	std::cout << "示例:" << std::endl;
	std::cout << "  " << programName << "                    # 签名所有文件" << std::endl;
	std::cout << "  " << programName << " document.txt       # 签名指定文件" << std::endl;
}

/**
 * @brief 主函数
 */
int main(int argc, char *argv[])
{
	const char *privateKeyPath = "private_key.pem";
	std::vector<std::string> filesToSign;

	std::cout << "=== XNSignature 文件签名工具 ===" << std::endl;

	if (argc == 1) {
		// 没有参数，签名目录下所有文件
		std::cout << "未指定文件，将签名当前目录下的所有文件..." << std::endl;
		filesToSign = getFilesInDirectory(".");

		if (filesToSign.empty()) {
			std::cout << "当前目录下没有找到可签名的文件" << std::endl;
			return 0;
		}

		std::cout << "找到 " << filesToSign.size() << " 个文件:" << std::endl;
		for (const auto &file : filesToSign) {
			std::cout << "  - " << file << std::endl;
		}
	} else if (argc == 2) {
		// 指定了单个文件
		filesToSign.push_back(argv[1]);
		std::cout << "将签名指定文件: " << argv[1] << std::endl;
	} else {
		std::cerr << "错误：参数数量不正确" << std::endl;
		showUsage(argv[0]);
		return -1;
	}

	std::cout << "私钥文件: " << privateKeyPath << std::endl;
	std::cout << "正在签名文件..." << std::endl;

	int successCount = 0;
	int totalCount	 = filesToSign.size();

	for (const auto &filePath : filesToSign) {
		// 生成签名文件路径
		std::string signaturePath = filePath + ".sig";

		std::cout << "\n正在签名: " << filePath << std::endl;
		std::cout << "签名文件: " << signaturePath << std::endl;

		int result = signFile(filePath.c_str(), privateKeyPath, signaturePath.c_str());

		if (result == 0) {
			std::cout << "✓ 签名成功" << std::endl;
			successCount++;
		} else {
			std::cout << "✗ 签名失败" << std::endl;
		}
	}

	std::cout << "\n=== 签名完成 ===" << std::endl;
	std::cout << "成功签名: " << successCount << "/" << totalCount << " 个文件" << std::endl;

	if (successCount == totalCount) {
		std::cout << "所有文件签名完成！" << std::endl;
		return 0;
	} else {
		std::cerr << "部分文件签名失败！" << std::endl;
		return -1;
	}
}