#include "pch.h"
#include "SSLClientConnection.h"

namespace SSLClient {

SSLClientConnection::SSLClientConnection()
	: m_sslContext(nullptr)
	, m_ssl(nullptr)
	, m_socket(INVALID_SOCKET)
	, m_isConnected(false)
{
}

SSLClientConnection::~SSLClientConnection()
{
	Cleanup();
}

bool SSLClientConnection::InitializeWinsock()
{
	WSADATA wsaData;
	int result = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if (result != 0) {
		std::cerr << "[错误] WSAStartup失败，错误码: " << result << std::endl;
		return false;
	}
	return true;
}

bool SSLClientConnection::Initialize(const char* clientCert, const char* clientKey,
                                      const char* caCert, const char* keyPassword)
{
	// 初始化Winsock
	if (!InitializeWinsock()) {
		return false;
	}

	std::cout << "[客户端] 正在初始化SSL环境..." << std::endl;

	// 创建SSL上下文
	m_sslContext = SSL_CTX_new(TLS_client_method());
	if (!m_sslContext) {
		CertificateManager::PrintSSLError("创建SSL上下文失败");
		return false;
	}

	// 设置验证模式：验证服务器证书
	SSL_CTX_set_verify(m_sslContext, SSL_VERIFY_PEER, nullptr);

	// 加载CA证书
	X509* ca = m_certManager.LoadCertificateFromMemory(caCert);
	if (!ca) {
		return false;
	}

	X509_STORE* store = SSL_CTX_get_cert_store(m_sslContext);
	if (X509_STORE_add_cert(store, ca) != 1) {
		CertificateManager::PrintSSLError("添加CA证书失败");
		X509_free(ca);
		return false;
	}
	X509_free(ca);

	// 加载客户端证书
	X509* cert = m_certManager.LoadCertificateFromMemory(clientCert);
	if (!cert) {
		return false;
	}

	if (SSL_CTX_use_certificate(m_sslContext, cert) != 1) {
		CertificateManager::PrintSSLError("使用客户端证书失败");
		X509_free(cert);
		return false;
	}
	X509_free(cert);

	// 加载客户端私钥
	EVP_PKEY* pkey = m_certManager.LoadPrivateKeyFromMemory(clientKey, keyPassword);
	if (!pkey) {
		return false;
	}

	if (SSL_CTX_use_PrivateKey(m_sslContext, pkey) != 1) {
		CertificateManager::PrintSSLError("使用客户端私钥失败");
		EVP_PKEY_free(pkey);
		return false;
	}
	EVP_PKEY_free(pkey);

	// 验证私钥和证书是否匹配
	if (SSL_CTX_check_private_key(m_sslContext) != 1) {
		std::cerr << "[错误] 客户端私钥和证书不匹配" << std::endl;
		return false;
	}

	std::cout << "[客户端] SSL环境初始化成功" << std::endl;
	return true;
}

bool SSLClientConnection::Connect(const char* address, int port)
{
	// 创建socket
	m_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (m_socket == INVALID_SOCKET) {
		std::cerr << "[错误] 创建socket失败，错误码: " << WSAGetLastError() << std::endl;
		return false;
	}

	// 设置服务器地址
	sockaddr_in serverAddr;
	serverAddr.sin_family = AF_INET;
	serverAddr.sin_port = htons(static_cast<u_short>(port));
	inet_pton(AF_INET, address, &serverAddr.sin_addr);

	// 连接到服务器
	std::cout << "[客户端] 正在连接服务器 " << address << ":" << port << " ..." << std::endl;
	if (connect(m_socket, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
		std::cerr << "[错误] 连接失败，错误码: " << WSAGetLastError() << std::endl;
		closesocket(m_socket);
		m_socket = INVALID_SOCKET;
		return false;
	}

	std::cout << "[客户端] TCP连接成功" << std::endl;

	// 创建SSL对象
	m_ssl = SSL_new(m_sslContext);
	if (!m_ssl) {
		CertificateManager::PrintSSLError("创建SSL对象失败");
		closesocket(m_socket);
		m_socket = INVALID_SOCKET;
		return false;
	}

	// 将SSL绑定到socket
	if (SSL_set_fd(m_ssl, static_cast<int>(m_socket)) != 1) {
		CertificateManager::PrintSSLError("绑定SSL到socket失败");
		SSL_free(m_ssl);
		m_ssl = nullptr;
		closesocket(m_socket);
		m_socket = INVALID_SOCKET;
		return false;
	}

	// 执行SSL握手
	std::cout << "[客户端] 正在进行SSL握手..." << std::endl;
	int ret = SSL_connect(m_ssl);
	if (ret != 1) {
		int err = SSL_get_error(m_ssl, ret);
		std::cerr << "[错误] SSL握手失败，错误码: " << err << std::endl;
		CertificateManager::PrintSSLError("SSL_connect");
		SSL_free(m_ssl);
		m_ssl = nullptr;
		closesocket(m_socket);
		m_socket = INVALID_SOCKET;
		return false;
	}

	m_isConnected = true;
	std::cout << "[客户端] SSL握手完成！" << std::endl;
	std::cout << "[客户端] 使用的加密套件: " << GetCipherSuite() << std::endl;

	return true;
}

bool SSLClientConnection::Send(const std::string& data)
{
	if (!m_isConnected || !m_ssl) {
		std::cerr << "[错误] 未连接到服务器" << std::endl;
		return false;
	}

	int sent = SSL_write(m_ssl, data.c_str(), static_cast<int>(data.length()));
	if (sent <= 0) {
		int err = SSL_get_error(m_ssl, sent);
		std::cerr << "[错误] 发送失败，错误码: " << err << std::endl;
		CertificateManager::PrintSSLError("SSL_write");
		return false;
	}

	std::cout << "[客户端] 发送数据: \"" << data << "\" (" << sent << " 字节)" << std::endl;
	return true;
}

int SSLClientConnection::Receive(char* buffer, int bufferSize)
{
	if (!m_isConnected || !m_ssl || !buffer || bufferSize <= 0) {
		return -1;
	}

	// 设置socket为非阻塞模式
	u_long mode = 1;
	ioctlsocket(m_socket, FIONBIO, &mode);

	int received = SSL_read(m_ssl, buffer, bufferSize - 1);

	// 恢复为阻塞模式
	mode = 0;
	ioctlsocket(m_socket, FIONBIO, &mode);

	if (received > 0) {
		buffer[received] = '\0';
		std::cout << "[客户端] 收到数据: \"" << buffer << "\" (" << received << " 字节)" << std::endl;
		return received;
	}
	else if (received < 0) {
		int err = SSL_get_error(m_ssl, received);
		// SSL_ERROR_WANT_READ 和 SSL_ERROR_WANT_WRITE 表示没有数据，不是错误
		if (err != SSL_ERROR_WANT_READ && err != SSL_ERROR_WANT_WRITE) {
			if (err != SSL_ERROR_ZERO_RETURN) { // 忽略正常关闭
				std::cerr << "[错误] 接收失败，错误码: " << err << std::endl;
			}
			return -1;
		}
	}

	return 0; // 没有数据
}

void SSLClientConnection::Disconnect()
{
	std::cout << "[客户端] 正在断开连接..." << std::endl;
	Cleanup();
	m_isConnected = false;
	std::cout << "[客户端] 已断开连接" << std::endl;
}

const char* SSLClientConnection::GetCipherSuite() const
{
	if (m_ssl) {
		return SSL_get_cipher(m_ssl);
	}
	return "未知";
}

void SSLClientConnection::Cleanup()
{
	if (m_ssl) {
		SSL_shutdown(m_ssl);
		SSL_free(m_ssl);
		m_ssl = nullptr;
	}

	if (m_socket != INVALID_SOCKET) {
		closesocket(m_socket);
		m_socket = INVALID_SOCKET;
	}

	if (m_sslContext) {
		SSL_CTX_free(m_sslContext);
		m_sslContext = nullptr;
	}

	WSACleanup();
}

} // namespace SSLClient