#include "SimpleSock.h"
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#pragma comment(lib, "ws2_32.lib")

namespace SimpleSock {

	void FetchLastError(int& code, std::string& msg)
	{
		code = GetLastError();
		LPSTR buffer = nullptr;
		size_t size = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS, NULL, code, 0, (LPSTR)&buffer, 0, NULL);
		msg = std::string(buffer, size);
		LocalFree(buffer);
	}

	sock::sock()
	{
		fd = -1;
		_code = 0;
	}

	sock::sock(sock&& s) noexcept
	{
		fd = s.fd;
		s.fd = -1;
		_code = 0;
	}

	sock::~sock()
	{
		close();
	}

	int sock::create_socket()
	{
		fd = socket(AF_INET, SOCK_STREAM, 0);
		if (fd < 0)
		{
			FetchLastError(_code, _msg);
			return -1;
		}
		return 0;
	}

	int sock::connect(const std::string& ip, int port)
	{
		if (fd > 0)
		{
			_code = -1;
			_msg = "[SimpleSock] Socket already connected.";
			return -2;
		}

		if (create_socket() < 0) return -1;

		struct sockaddr_in saddr;
		memset(&saddr, 0, sizeof(saddr));
		if (inet_pton(AF_INET, ip.c_str(), &(saddr.sin_addr.s_addr)) != 1)
		{
			FetchLastError(_code, _msg);
			return -1;
		}
		saddr.sin_port = htons(port);
		saddr.sin_family = AF_INET;

		int ret = ::connect(fd, (sockaddr*)&saddr, sizeof(saddr));
		if (ret < 0)
		{
			FetchLastError(_code, _msg);
			return -1;
		}

		return 0;
	}

	int sock::bind(const std::string& ip, int port)
	{
		if (fd > 0)
		{
			_code = -1;
			_msg = "[SimpleSock] Socket already binded.";
			return -1;
		}

		struct sockaddr_in saddr;
		memset(&saddr, 0, sizeof(saddr));

		if (inet_pton(AF_INET, ip.c_str(), &(saddr.sin_addr.s_addr)) != 1)
		{
			FetchLastError(_code, _msg);
			return -1;
		}
		saddr.sin_port = htons(port);
		saddr.sin_family = AF_INET;

		int ret = ::bind(fd, (sockaddr*)&saddr, sizeof(saddr));
		if (ret < 0)
		{
			FetchLastError(_code, _msg);
			return -1;
		}

		return 0;
	}

	int sock::listen(int backlog)
	{
		int ret = ::listen(fd, backlog);
		if (ret < 0)
		{
			FetchLastError(_code, _msg);
			return -1;
		}
		return 0;
	}

	sock sock::accept()
	{
		sockaddr_in saddr;
		socklen_t saddrsz = sizeof(saddr);
		int ret = ::accept(fd, (sockaddr*)&saddr, &saddrsz);
		if (ret < 0)
		{
			FetchLastError(_code, _msg);
			return sock();
		}
		else
		{
			sock s;
			s.fd = ret;
			return s;
		}
	}

	int sock::send(const void* buffer, size_t length)
	{
		int ret = ::send(fd, (char*)buffer, length, 0);
		if (ret <= 0)
		{
			FetchLastError(_code, _msg);
		}
		return ret;
	}

	int sock::sendall(const void* buffer, size_t length)
	{
		size_t done = 0;
		while (done < length)
		{
			int ret = send((char*)buffer + done, length - done);
			if (ret <= 0) return ret;
			done += ret;
		}
		return done;
	}

	int sock::recv(void* buffer, size_t length)
	{
		int ret = ::recv(fd, (char*)buffer, length, 0);
		if (ret <= 0)
		{
			FetchLastError(_code, _msg);
		}
		return ret;
	}

	void sock::close()
	{
		if (fd > 0)
		{
			closesocket(fd);
			fd = -1;
		}
	}

	bool sock::valid()
	{
		return fd >= 0;
	}

	int sock::getError()
	{
		return _code;
	}

	std::string sock::getErrorMessage()
	{
		return _msg;
	}

	std::tuple<std::vector<std::string>, int, std::string> DNSResolve(const std::string& host)
	{
		std::vector<std::string> vec;
		int _code;
		std::string _msg;

		struct addrinfo hints;
		memset(&hints, 0, sizeof(hints));
		hints.ai_family = AF_UNSPEC;
		hints.ai_socktype = SOCK_STREAM;
		hints.ai_protocol = IPPROTO_TCP;

		struct addrinfo* result = nullptr;

		int ret = getaddrinfo(host.c_str(), NULL, &hints, &result);
		if (ret)
		{
			FetchLastError(_code, _msg);
			return std::make_tuple(vec, _code, _msg);
		}

		for (struct addrinfo* ptr = result; ptr; ptr = ptr->ai_next)
		{
			if (ptr->ai_family == AF_INET)
			{
				struct sockaddr_in* paddr = (struct sockaddr_in*)(ptr->ai_addr);
				char ip_buff[128] = { 0 };
				const char* ptr = inet_ntop(AF_INET, &(paddr->sin_addr), ip_buff, 128);
				if (ptr)
				{
					vec.push_back(ptr);
				}
				else
				{
					FetchLastError(_code, _msg);
				}
			}
		}

		freeaddrinfo(result);

		return std::make_tuple(vec, _code, _msg);
	}

	int ConnectWithHttpProxy(sock& s, const std::string& proxyIP, int proxyPort, const std::string& host, int port)
	{
		int ret = s.connect(proxyIP, proxyPort);
		if (ret < 0) return -1;
		char request[1024] = { 0 };
		snprintf(request, sizeof(request), "CONNECT %s:%d HTTP/1.1\r\nHost: %s:%d\r\n\r\n", host.c_str(), port, host.c_str(), port);
		ret = s.sendall(request, strlen(request));
		if (ret != strlen(request)) return -2;
		char buffer[8] = { 0 };
		std::string result;
		while (1)
		{
			ret = s.recv(buffer, 1);
			if (ret <= 0) return ret;
			result.push_back(buffer[0]);
			if (result.find("\r\n\r\n") != std::string::npos)
			{
				printf("%s\n===== HTTP PROXY CONNECT FINISHED =====\n", result.c_str());
				break;
			}
		}
		return 0;
	}

	int InitSimpleSock()
	{
		WORD wd;
		WSAData wdt;
		wd = MAKEWORD(2, 2);
		return WSAStartup(wd, &wdt);
	}
}