#include "SocketS.h"
#include <iostream>
#define DOSTART do{
#define DOEND } while (0);

//工作者线程
DWORD WINAPI SocketS::ThreadProc(LPVOID lpParameter)
{
	SocketS* server = (SocketS*)lpParameter;
	HANDLE port = server->hPort;
	DWORD len;
	ULONG_PTR index;
	LPOVERLAPPED lpOverlapped;

	while (server->g_flag)
	{//获取完成端口的数据
		BOOL bFlag = GetQueuedCompletionStatus(port, &len, &index, &lpOverlapped, INFINITE);
		clienInfo* clien = &server->ClienMap[index];
		if (bFlag == TRUE && len > 0)
		{// 接收到客户端消息
			server->Lock.lock();
			if (server->m_mode == 0)
			{//普通模式
				if (server->Rfunc != NULL)
				{
					server->Rfunc(index, clien->buff, len);
				}
				ZeroMemory(clien->buff, sizeof(clien->buff));
				server->PostRecv(index);	// 对自己投递接收
			}
			else if (server->m_mode == 1)
			{
				clien->Rbuflen += len;
				DOSTART
					if (clien->Rbuflen < sizeof(MsgHead))
						break;
				if (clien->Rbuflen - sizeof(MsgHead) < clien->h.bufLen)
					break;
				if (clien->buff == NULL)
					break;
				if (server->Rfunc == NULL)
					break;
				if (clien->h.isStr == 1)
				{
					AES aes;
					string str = aes.DecryptionAES(clien->buff, "qwertyuiopasdfgh");
					server->Rfunc(index, (char*)str.c_str(), str.length());
				}
				else
					server->Rfunc(index, clien->buff, len);

				ZeroMemory(clien->buff, sizeof(clien->buff));
				clien->Rbuflen = 0;
				DOEND
					server->PostRecv(index);	// 对自己投递接收
			}
			server->Lock.unlock();
		}
		else
		{//客户端关闭
			server->Lock.lock();
			if (server->Cfunc != NULL)
			{
				server->Cfunc(index);
			}
			server->CloseClien(index);
			server->Lock.unlock();
		}

	}
	server->Count--;
	return 0;
}

//构造
SocketS::SocketS(int mode)
{
	m_mode = mode;
}

//析构
SocketS::~SocketS()
{
	Close();
}

//投递消息
VOID SocketS::PostRecv(int index)
{
	WSABUF wsabuf{};
	clienInfo* clien = &ClienMap[index];
	if (m_mode == 0)
	{
		wsabuf.buf = clien->buff;
		wsabuf.len = 1020;
	}
	else if (m_mode == 1)
	{
		if (clien->Rbuflen < sizeof(MsgHead))
		{//收到的消息还没消息头长
			wsabuf.buf = (char*)&clien->h + clien->Rbuflen;
			wsabuf.len = sizeof(MsgHead) - clien->Rbuflen;
		}
		else
		{
			if (clien->Rbuflen == sizeof(MsgHead))
			{//刚接收完消息头
				if (!CheckHead(&clien->h))
				{//消息头有问题
					CloseClien(index);
					return;
				}
			}
			wsabuf.buf = clien->buff + (clien->Rbuflen - sizeof(MsgHead));
			//剩余消息如果大于1020那么分块接收消息
			if (clien->h.bufLen - clien->Rbuflen + sizeof(MsgHead) > 1020)
				wsabuf.len = 1020;
			else
				wsabuf.len = clien->h.bufLen - clien->Rbuflen + sizeof(MsgHead);
		}
	}

	//让完成端口继续接收数据
	DWORD dwRecvCount = wsabuf.len;
	DWORD dwFlag = 0;
	WSARecv(clien->c_Sock, &wsabuf, 1, &dwRecvCount, &dwFlag, &clien->c_Olp, NULL);

	//if (ERROR_IO_PENDING != WSAGetLastError())
	//{//消息投递失败
	//	return;	
	//}
}

//给指定客户端发送数据
int SocketS::PostSend(int index, const char* buf, int len)
{
	//判断客户端存不存在
	if (ClienMap.find(index) != ClienMap.end())
	{
		if (ClienMap[index].c_Sock != 0)
		{
			return send(ClienMap[index].c_Sock, buf, len, 0);
		}
	}
	return -1;
}

//检查消息头
BOOL SocketS::CheckHead(MsgHead* h)
{
	//消息长度不能大于10000
	if (h->bufLen > 10000)
	{
		return FALSE;
	}
	//机密数据并验证发过来的Token是否正确
	string str, str1;
	str = to_string(h->tm);
	str += to_string(h->bufLen);
	md5 md;
	str1 = md.StringToMD5(str);
	str = h->token;
	if (str == str1)
	{//token正确验证发送消息的时间
		time_t t1;
		time(&t1);
		if (abs(t1 - h->tm) < 300)
		{
			return TRUE;
		}
	}
	return FALSE;
}

//监听消息函数
VOID SocketS::Listens()
{
	SOCKADDR_IN client;
	int iaddrSize = sizeof(SOCKADDR_IN);
	while (TRUE)
	{//开始监听并绑定客户端
		ClienMap[g_count].c_Sock = accept(ClienMap[0].c_Sock, (struct sockaddr*)&client, &iaddrSize);   
		if (INVALID_SOCKET == ClienMap[g_count].c_Sock)
		{//客户端绑定失败
			return;
		}
		HANDLE hPort1 = CreateIoCompletionPort((HANDLE)ClienMap[g_count].c_Sock, hPort, g_count, 0);
		if (hPort1 != hPort)
		{//完成端口绑定失败
			printf("CreateIoCompletionPort 绑定失败 error:%d\n", GetLastError());
			closesocket(ClienMap[g_count].c_Sock);
			continue;
		}
		PostRecv(g_count);	// 投递recv
		std::cout << g_count << "连接成功\n";
		++g_count;			// 客户端数量++
	}
}

//创建服务端
BOOL SocketS::Creat(int Prot)
{
	if (g_flag == 1)
		return FALSE;
	//创建服务端
	g_count = 1;
	SOCKADDR_IN local;
	WSADATA wsaData;
	DWORD dwThreadId;
	if (WSAStartup(0x0202, &wsaData) != 0)
	{
		return FALSE;
	}
	//创建完成端口
	hPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
	//获取CPU核心数，并准备创先线程
	SYSTEM_INFO systemProcessorsCount;
	GetSystemInfo(&systemProcessorsCount);
	Count = systemProcessorsCount.dwNumberOfProcessors;
	g_flag = 1;
	//创建线程
	for (int i = 0; i < Count; i++)
	{
		CreateThread(NULL, 0, ThreadProc, (LPVOID)this, 0, &dwThreadId);
	}

	//开始绑定服务端
	ClienMap[0].c_Sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	local.sin_addr.S_un.S_addr = htonl(INADDR_ANY);
	local.sin_family = AF_INET;
	local.sin_port = htons(Prot);
	if (::bind(ClienMap[0].c_Sock, (struct sockaddr*)&local, sizeof(SOCKADDR_IN)) != 0)
	{//绑定失败
		g_flag = 0;
		PostQueuedCompletionStatus(hPort, 0xFFFFFFFF, 0, NULL);
		CloseHandle(hPort);
		closesocket(ClienMap[0].c_Sock);
		WSACleanup();
		return FALSE;
	}
	//创建线程监听
	listen(ClienMap[0].c_Sock, 5);
	ListenClien = new std::thread((&SocketS::Listens), this);
	return TRUE;
}

//关闭服务端
void SocketS::Close()
{//清除客户端
	if (g_flag == 0)
		return;
	g_flag = 0;
	Lock.lock();
	for (auto i = ClienMap.begin(); i != ClienMap.end(); i++)
	{
		closesocket(i->second.c_Sock);
	}
	ClienMap.clear();
	Lock.unlock();
	//关闭工作者线程并等待所有线程退出
	ListenClien->join();
	delete ListenClien;
	ListenClien = NULL;
	CloseHandle(hPort);
	while (Count)
	{
		Sleep(50);
	}
	WSACleanup();
}

//设置接收消息的回调函数
VOID SocketS::SetRecvFunc(function<VOID(int, char*, int)> fun)
{
	Rfunc = fun;
}

//设置关闭客户端的回调函数
VOID SocketS::SetCloseFunc(function<VOID(int)> fun)
{
	Cfunc = fun;
}

//获取客户端IP和端口
VOID SocketS::GetClienName(int index, string& IP, int& Prot)
{
	IP = "";
	Prot = 0;
	if(ClienMap.find(index)!= ClienMap.end())
	{
		if(ClienMap[index].c_Sock!=0)
		{
			struct sockaddr_in peerAddr;
			int peerLen = sizeof(peerAddr);
			getpeername(ClienMap[index].c_Sock, (struct sockaddr*)&peerAddr, &peerLen);
			IP = inet_ntoa(peerAddr.sin_addr);
			Prot = ntohs(peerAddr.sin_port);
		}
	}
}

//发送数据包
int SocketS::SendData(int index,const char* buf, int len, int isStr)
{
	if (buf == NULL)
		return -1;
	if (m_mode == 0)
	{//普通模式
		return PostSend(index, buf, len);
	}
	else if (m_mode == 1)
	{
		//创建消息头生成token
		MsgHead h;
		h.bufLen = len;
		h.isStr = isStr;
		time(&h.tm);
		string str = to_string(h.tm);
		str += to_string(len);
		md5 md;
		str = md.StringToMD5(str);
		memcpy_s(h.token, str.length(), str.c_str(), str.length());
		//合并消息头和消息体
		char* buff = new char[sizeof(MsgHead) + len];
		memcpy_s(buff, sizeof(MsgHead), &h, sizeof(MsgHead));
		memcpy_s(buff + sizeof(MsgHead), len, buf, len);
		//发送消息
		int lenth = PostSend(index, buff, len + sizeof(MsgHead)) - sizeof(MsgHead);
		delete[] buff;
		if (lenth < 0)
		{//消息发送失败
			CloseClien(index);
			lenth = -1;
		}
		return lenth;
	}
	return -1;
}

//发送字符串
BOOL SocketS::SendStr(int index, string str)
{//判断消息不能为空
	if (str == "")
		return FALSE;
	if (m_mode == 0)
	{//普通模式
		if (PostSend(index, str.c_str(), str.length()) == -1)
		{
			return FALSE;
		}
	}
	else if (m_mode == 1)
	{
		//加密字符串
		AES aes;
		string sendStr = aes.EncryptionAES(str, "qwertyuiopasdfgh");
		//发送消息
		if (-1 == SendData(index, sendStr.c_str(), sendStr.length(), 1))
		{//消息发送失败
			return FALSE;
		}
	}
	return TRUE;
}

//关闭指定客户端
VOID SocketS::CloseClien(int index)
{
	if (ClienMap.find(index) != ClienMap.end())
	{//客户端存在断开连接
		cout << index << "断开连接\n";
		shutdown(ClienMap[index].c_Sock, SD_BOTH);
		closesocket(ClienMap[index].c_Sock);
		ClienMap[index].c_Sock = 0;
		ClienMap[index].c_Olp.hEvent = NULL;
		ClienMap.erase(index);
	}
}