#include "UdpSocket.h"

void UDPServer::cleanup() {
#ifdef _WIN32
	if (serverSocket != INVALID_SOCKET) {
		closesocket(serverSocket);
		WSACleanup();
	}
#else
	if (serverSocket != INVALID_SOCKET) {
		close(serverSocket);
	}
#endif
}

UDPServer::UDPServer() : serverSocket(INVALID_SOCKET), initialized(false) {
#ifdef _WIN32
	WSADATA wsaData;
	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
		std::cerr << "WSAStartup failed\n";
		return;
	}
#endif
	initialized = true;
}

UDPServer::~UDPServer() {
	cleanup();
}

bool UDPServer::Start(int port) {
	if (!initialized) return false;

	serverSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (serverSocket == INVALID_SOCKET) {
		std::cerr << "Create socket failed\n";
		return false;
	}

	serverAddr.sin_family = AF_INET;
	serverAddr.sin_port = htons(port);
	serverAddr.sin_addr.s_addr = htonl(INADDR_ANY);

	if (bind(serverSocket, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
		std::cerr << "Bind failed\n";
		return false;
	}

	return true;
}

// 修改后的接收方法（UDPServer类）
int UDPServer::Receive(char* buffer, int bufferSize, std::string& clientIP, int& clientPort) {
	sockaddr_in clientAddr;
	socklen_t clientAddrLen = sizeof(clientAddr);

	int recvSize = recvfrom(serverSocket, buffer, bufferSize, 0,
		(sockaddr*)&clientAddr, &clientAddrLen);
	if (recvSize == SOCKET_ERROR) {
		return -1;
	}

	// 使用跨平台的地址转换方法
	char ipBuffer[INET_ADDRSTRLEN] = { 0 };
#ifdef _WIN32
	InetNtopA(AF_INET, &clientAddr.sin_addr, ipBuffer, INET_ADDRSTRLEN);
#else
	inet_ntop(AF_INET, &clientAddr.sin_addr, ipBuffer, INET_ADDRSTRLEN);
#endif

	clientIP = ipBuffer;
	clientPort = ntohs(clientAddr.sin_port);
	return recvSize;
}

bool UDPServer::Send(const char* buffer, int bufferSize, const std::string& clientIP, int clientPort) {
	sockaddr_in clientAddr;
	clientAddr.sin_family = AF_INET;
	clientAddr.sin_port = htons(clientPort);
	inet_pton(AF_INET, clientIP.c_str(), &clientAddr.sin_addr);

	int sentSize = sendto(serverSocket, buffer, bufferSize, 0,
		(sockaddr*)&clientAddr, sizeof(clientAddr));
	return sentSize == bufferSize;
}

void UDPServer::Close() {
	cleanup();
	serverSocket = INVALID_SOCKET;
}


/// <summary>
/// UDPClient
/// </summary>
void UDPClient::cleanup() {
#ifdef _WIN32
	if (clientSocket != INVALID_SOCKET) {
		closesocket(clientSocket);
		WSACleanup();
	}
#else
	if (clientSocket != INVALID_SOCKET) {
		close(clientSocket);
	}
#endif
}

UDPClient::UDPClient() : clientSocket(INVALID_SOCKET), initialized(false) {
#ifdef _WIN32
	WSADATA wsaData;
	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
		std::cerr << "WSAStartup failed\n";
		return;
	}
#endif
	initialized = true;
}

UDPClient::~UDPClient() {
	cleanup();
}

bool UDPClient::Connect(const std::string& serverIP, int serverPort) {
	if (!initialized) return false;

	clientSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (clientSocket == INVALID_SOCKET) {
		std::cerr << "Create socket failed\n";
		return false;
	}

	serverAddr.sin_family = AF_INET;
	serverAddr.sin_port = htons(serverPort);
	inet_pton(AF_INET, serverIP.c_str(), &serverAddr.sin_addr);

	return true;
}

bool UDPClient::Send(const char* buffer, int bufferSize) {
	int sentSize = sendto(clientSocket, buffer, bufferSize, 0,
		(sockaddr*)&serverAddr, sizeof(serverAddr));
	return sentSize == bufferSize;
}

//lTimeoutMs 单位毫秒
int UDPClient::Receive(char* buffer, int bufferSize, int lTimeoutMs)
{
	sockaddr_in fromAddr;
	socklen_t fromLen = sizeof(fromAddr);

#ifdef _WIN32
	int timeout = lTimeoutMs;  // 转换为毫秒
	setsockopt(clientSocket, SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout, sizeof(timeout));
#else
	struct timeval tv;
	tv.tv_sec = lTimeoutMs / 1000;
	tv.tv_usec = (lTimeoutMs % 1000) * 1000;
	setsockopt(clientSocket, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
#endif

	int recvSize = recvfrom(clientSocket, buffer, bufferSize, 0,
		(sockaddr*)&fromAddr, &fromLen);
	return recvSize;
}

// sbuffer 发送数据
// sbufferSize 发送数据长度
// rbuffer 接收数据缓冲区
// rbufferSize 接收数据长度
// lRto 重传时间
// lRtoCnt 重传次数
// 返回值 接收数据长度,小于等于零表示发送以及重传都失败了
int UDPClient::SendAndRecv(const char* sbuffer, int sbufferSize, char* rbuffer, int rbufferSize, int lRto, int lRtoCnt)
{
	do
	{
		if (!Send(sbuffer, sbufferSize))
		{
			continue;
		}

		int lRecvLen = Receive(rbuffer, rbufferSize, lRto);
		if (lRecvLen > 0)
		{
			return lRecvLen;
		}

	} while (lRtoCnt-- > 0);

	return -1;
}

void UDPClient::Close() {
	cleanup();
	clientSocket = INVALID_SOCKET;
}