#include <filezilla.h>
#include "externalipresolver.h"
#include "socket.h"
#include "misc.h"

#include <libfilezilla/format.hpp>
#include <libfilezilla/iputils.hpp>

#include <regex>

namespace {
fz::mutex s_sync;
std::string ip;
bool checked = false;
}

CExternalIPResolver::CExternalIPResolver(fz::thread_pool & pool, fz::event_handler & handler)
	: fz::event_handler(handler.event_loop_)
	, thread_pool_(pool)
	, m_handler(&handler)
{
	ResetHttpData(true);
}

CExternalIPResolver::~CExternalIPResolver()
{
	remove_handler();

	delete [] m_pRecvBuffer;
	m_pRecvBuffer = 0;

	delete socket_;
	socket_ = 0;
}

void CExternalIPResolver::GetExternalIP(std::wstring const& address, fz::address_type protocol, bool force)
{
	{
		fz::scoped_lock l(s_sync);
		if (checked) {
			if (force)
				checked = false;
			else {
				m_done = true;
				return;
			}
		}
	}

	m_address = address;
	m_protocol = protocol;

	std::wstring host;
	size_t pos = address.find(L"://");
	if (pos != std::wstring::npos) {
		host = address.substr(pos + 3);
	}
	else {
		host = address;
	}

	pos = host.find('/');
	if (pos != std::wstring::npos) {
		host = host.substr(0, pos);
	}

	std::wstring hostWithPort = host;
	pos = host.rfind(':');
	if (pos != std::wstring::npos) {
		std::wstring port = host.substr(pos + 1);
		m_port = fz::to_integral<decltype(m_port)>(port);
		if (m_port < 1 || m_port > 65535) {
			m_port = 80;
		}
		host = host.substr(0, pos);
	}
	else {
		m_port = 80;
	}

	if (host.empty()) {
		m_done = true;
		return;
	}

	socket_ = new fz::socket(thread_pool_, this);

	int res = socket_->connect(fz::to_native(host), m_port, protocol);
	if (res && res != EINPROGRESS) {
		Close(false);
		return;
	}

	m_sendBuffer = fz::sprintf("GET %s HTTP/1.1\r\nHost: %s\r\nUser-Agent: %s\r\nConnection: close\r\n\r\n", fz::to_utf8(address), fz::to_utf8(hostWithPort), fz::replaced_substrings(PACKAGE_STRING, " ", "/"));
}

void CExternalIPResolver::operator()(fz::event_base const& ev)
{
	fz::dispatch<fz::socket_event>(ev, this, &CExternalIPResolver::OnSocketEvent);
}

void CExternalIPResolver::OnSocketEvent(fz::socket_event_source*, fz::socket_event_flag t, int error)
{
	if (!socket_) {
		return;
	}

	switch (t)
	{
	case fz::socket_event_flag::read:
		OnReceive();
		break;
	case fz::socket_event_flag::connection:
		OnConnect(error);
		break;
	case fz::socket_event_flag::close:
		OnClose();
		break;
	case fz::socket_event_flag::write:
		OnSend();
		break;
	default:
		break;
	}

}

void CExternalIPResolver::OnConnect(int error)
{
	if (error)
		Close(false);
}

void CExternalIPResolver::OnClose()
{
	if (m_data.empty())
		Close(false);
	else
		OnData(0, 0);
}

void CExternalIPResolver::OnReceive()
{
	if (!m_pRecvBuffer) {
		m_pRecvBuffer = new char[m_recvBufferLen];
		m_recvBufferPos = 0;
	}

	if (!m_sendBuffer.empty()) {
		return;
	}

	while (socket_) {
		unsigned int len = m_recvBufferLen - m_recvBufferPos;
		int error;
		int read = socket_->read(m_pRecvBuffer + m_recvBufferPos, len, error);
		if (read == -1) {
			if (error != EAGAIN)
				Close(false);
			return;
		}

		if (!read) {
			Close(false);
			return;
		}

		if (m_finished) {
			// Just ignore all further data
			m_recvBufferPos = 0;
			return;
		}

		m_recvBufferPos += read;

		if (!m_gotHeader)
			OnHeader();
		else {
			if (m_transferEncoding == chunked)
				OnChunkedData();
			else
				OnData(m_pRecvBuffer, m_recvBufferPos);
		}
	}
}

void CExternalIPResolver::OnSend()
{
	while (!m_sendBuffer.empty()) {
		int error;
		int written = socket_->write(m_sendBuffer.c_str(), m_sendBuffer.size(), error);
		if (written == -1) {
			if (error != EAGAIN) {
				Close(false);
			}
			return;
		}

		if (!written) {
			Close(false);
			return;
		}

		m_sendBuffer = m_sendBuffer.substr(written);
		if (m_sendBuffer.empty()) {
			OnReceive();
		}
	}
}

void CExternalIPResolver::Close(bool successful)
{
	m_sendBuffer.clear();

	delete [] m_pRecvBuffer;
	m_pRecvBuffer = 0;

	delete socket_;
	socket_ = 0;

	if (m_done) {
		return;
	}

	m_done = true;

	{
		fz::scoped_lock l(s_sync);
		if (!successful) {
			ip.clear();
		}
		checked = true;
	}

	if (m_handler) {
		m_handler->send_event<CExternalIPResolveEvent>();
		m_handler = 0;
	}
}

void CExternalIPResolver::OnHeader()
{
	// Parse the HTTP header.
	// We do just the neccessary parsing and silently ignore most header fields
	// Redirects are supported though if the server sends the Location field.

	for (;;) {
		// Find line ending
		unsigned int i = 0;
		for (i = 0; (i + 1) < m_recvBufferPos; ++i) {
			if (m_pRecvBuffer[i] == '\r') {
				if (m_pRecvBuffer[i + 1] != '\n') {
					Close(false);
					return;
				}
				break;
			}
		}
		if ((i + 1) >= m_recvBufferPos) {
			if (m_recvBufferPos == m_recvBufferLen) {
				// We don't support header lines larger than 4096
				Close(false);
				return;
			}
			return;
		}

		m_pRecvBuffer[i] = 0;

		if (!m_responseCode) {
			m_responseString = m_pRecvBuffer;
			if (m_recvBufferPos < 16 || memcmp(m_pRecvBuffer, "HTTP/1.", 7)) {
				// Invalid HTTP Status-Line
				Close(false);
				return;
			}

			if (m_pRecvBuffer[9] < '1' || m_pRecvBuffer[9] > '5' ||
				m_pRecvBuffer[10] < '0' || m_pRecvBuffer[10] > '9' ||
				m_pRecvBuffer[11] < '0' || m_pRecvBuffer[11] > '9')
			{
				// Invalid response code
				Close(false);
				return;
			}

			m_responseCode = (m_pRecvBuffer[9] - '0') * 100 + (m_pRecvBuffer[10] - '0') * 10 + m_pRecvBuffer[11] - '0';

			if (m_responseCode >= 400) {
				// Failed request
				Close(false);
				return;
			}

			if (m_responseCode == 305) {
				// Unsupported redirect
				Close(false);
				return;
			}
		}
		else {
			if (!i) {
				// End of header, data from now on

				// Redirect if neccessary
				if (m_responseCode >= 300) {
					delete socket_;
					socket_ = 0;

					delete [] m_pRecvBuffer;
					m_pRecvBuffer = 0;

					std::wstring location = m_location;

					ResetHttpData(false);

					GetExternalIP(location, m_protocol);
					return;
				}

				m_gotHeader = true;

				memmove(m_pRecvBuffer, m_pRecvBuffer + 2, m_recvBufferPos - 2);
				m_recvBufferPos -= 2;
				if (m_recvBufferPos) {
					if (m_transferEncoding == chunked)
						OnChunkedData();
					else
						OnData(m_pRecvBuffer, m_recvBufferPos);
				}
				return;
			}
			if (m_recvBufferPos > 12 && !memcmp(m_pRecvBuffer, "Location: ", 10)) {
				m_location = fz::to_wstring_from_utf8(m_pRecvBuffer + 10);
			}
			else if (m_recvBufferPos > 21 && !memcmp(m_pRecvBuffer, "Transfer-Encoding: ", 19)) {
				if (!strcmp(m_pRecvBuffer + 19, "chunked"))
					m_transferEncoding = chunked;
				else if (!strcmp(m_pRecvBuffer + 19, "identity")) {
					m_transferEncoding = identity;
				}
				else {
					m_transferEncoding = unknown;
				}
			}
		}

		memmove(m_pRecvBuffer, m_pRecvBuffer + i + 2, m_recvBufferPos - i - 2);
		m_recvBufferPos -= i + 2;

		if (!m_recvBufferPos)
			break;
	}
}

void CExternalIPResolver::OnData(char* buffer, unsigned int len)
{
	if (buffer) {
		unsigned int i;
		for (i = 0; i < len; ++i) {
			if (buffer[i] == '\r' || buffer[i] == '\n')
				break;
			if (buffer[i] & 0x80) {
				Close(false);
				return;
			}
		}

		if (i) {
			m_data.append(buffer, i);
		}

		if (i == len) {
			return;
		}
	}

	if (m_protocol == fz::address_type::ipv6) {
		if (!m_data.empty() && m_data[0] == '[') {
			if (m_data.back() != ']') {
				Close(false);
				return;
			}
			m_data = m_data.substr(1, m_data.size() - 2);
		}

		if (fz::get_ipv6_long_form(m_data).empty()) {
			Close(false);
			return;
		}

		fz::scoped_lock l(s_sync);
		ip = m_data;
	}
	else {
		// Validate ip address
		std::string const digit = "0*[0-9]{1,3}";
		char const* const dot = "\\.";
		std::string const exp = "(^|[^\\.[:digit:]])(" + digit + dot + digit + dot + digit + dot + digit + ")([^\\.[:digit:]]|$)";

		std::regex const regex(exp);
		std::smatch m;
		if (!std::regex_search(m_data, m, regex)) {
			Close(false);
			return;
		}

		fz::scoped_lock l(s_sync);
		ip = m[2];
	}

	Close(true);
}

void CExternalIPResolver::ResetHttpData(bool resetRedirectCount)
{
	m_gotHeader = false;
	m_location.clear();
	m_responseCode = 0;
	m_responseString.clear();
	if (resetRedirectCount) {
		m_redirectCount = 0;
	}

	m_transferEncoding = unknown;

	m_chunkData = t_chunkData();

	m_finished = false;
}

void CExternalIPResolver::OnChunkedData()
{
	char* p = m_pRecvBuffer;
	unsigned int len = m_recvBufferPos;

	for (;;) {
		if (m_chunkData.size != 0) {
			unsigned int dataLen = len;
			if (m_chunkData.size < len)
				dataLen = static_cast<unsigned int>(m_chunkData.size);
			OnData(p, dataLen);
			if (!m_pRecvBuffer)
				return;

			m_chunkData.size -= dataLen;
			p += dataLen;
			len -= dataLen;

			if (m_chunkData.size == 0)
				m_chunkData.terminateChunk = true;

			if (!len)
				break;
		}

		// Find line ending
		unsigned int i = 0;
		for (i = 0; (i + 1) < len; ++i) {
			if (p[i] == '\r') {
				if (p[i + 1] != '\n') {
					Close(false);
					return;
				}
				break;
			}
		}
		if ((i + 1) >= len) {
			if (len == m_recvBufferLen) {
				// We don't support lines larger than 4096
				Close(false);
				return;
			}
			break;
		}

		p[i] = 0;

		if (m_chunkData.terminateChunk) {
			if (i) {
				// Chunk has to end with CRLF
				Close(false);
				return;
			}
			m_chunkData.terminateChunk = false;
		}
		else if (m_chunkData.getTrailer) {
			if (!i) {
				m_finished = true;
				m_recvBufferPos = 0;
				return;
			}

			// Ignore the trailer
		}
		else {
			// Read chunk size
			char* q = p;
			while (*q) {
				if (*q >= '0' && *q <= '9') {
					m_chunkData.size *= 16;
					m_chunkData.size += *q - '0';
				}
				else if (*q >= 'A' && *q <= 'F') {
					m_chunkData.size *= 10;
					m_chunkData.size += *q - 'A' + 10;
				}
				else if (*q >= 'a' && *q <= 'f') {
					m_chunkData.size *= 10;
					m_chunkData.size += *q - 'a' + 10;
				}
				else if (*q == ';' || *q == ' ')
					break;
				else {
					// Invalid size
					Close(false);
					return;
				}
				q++;
			}
			if (m_chunkData.size == 0)
				m_chunkData.getTrailer = true;
		}

		p += i + 2;
		len -= i + 2;

		if (!len)
			break;
	}

	if (p != m_pRecvBuffer) {
		memmove(m_pRecvBuffer, p, len);
		m_recvBufferPos = len;
	}
}

bool CExternalIPResolver::Successful() const
{
	fz::scoped_lock l(s_sync);
	return !ip.empty();
}

std::string CExternalIPResolver::GetIP() const
{
	fz::scoped_lock l(s_sync);
	return ip;
}
