/*
 * SocketListener.cpp -- portable select() wrapper
 *
 * Copyright (c) 2013, 2014 David Demelier <markand@malikania.fr>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <algorithm>
#include <map>
#include <set>
#include <utility>
#include <vector>

#include "SocketListener.h"

/* --------------------------------------------------------
 * Select implementation
 * -------------------------------------------------------- */

namespace backend {

void Select::set(Socket &s, int direction)
{
	if (m_table.count(s.handle()) > 0) {
		m_table.at(s.handle()).second |= direction;
	} else {
		m_table.insert({s.handle(), {s, direction}});
	}

}

void Select::unset(Socket &s, int direction)
{
	if (m_table.count(s.handle()) != 0) {
		m_table.at(s.handle()).second &= ~(direction);

		// If no read, no write is requested, remove it
		if (m_table.at(s.handle()).second == 0) {
			m_table.erase(s.handle());
		}
	}
}

void Select::remove(Socket &sc)
{
	m_table.erase(sc.handle());
}

void Select::clear()
{
	m_table.clear();
}

SocketStatus Select::wait(int ms)
{
	auto result = waitMultiple(ms);

	if (result.size() == 0) {
		throw SocketError(SocketError::System, "select", "No socket found");
	}

	return result[0];
}

std::vector<SocketStatus> Select::waitMultiple(int ms)
{
	timeval maxwait, *towait;
	fd_set readset;
	fd_set writeset;

	FD_ZERO(&readset);
	FD_ZERO(&writeset);

	Socket::Handle max = 0;

	for (auto &s : m_table) {
		if (s.second.second & SocketListener::Read) {
			FD_SET(s.first, &readset);
		}
		if (s.second.second & SocketListener::Write) {
			FD_SET(s.first, &writeset);
		}

		if (s.first > max) {
			max = s.first;
		}
	}

	maxwait.tv_sec = 0;
	maxwait.tv_usec = ms * 1000;

	// Set to nullptr for infinite timeout.
	towait = (ms < 0) ? nullptr : &maxwait;

	auto error = ::select(max + 1, &readset, &writeset, nullptr, towait);
	if (error == Socket::Error) {
		throw SocketError(SocketError::System, "select");
	}
	if (error == 0) {
		throw SocketError(SocketError::Timeout, "select", "Timeout while listening");
	}

	std::vector<SocketStatus> sockets;

	for (auto &c : m_table) {
		if (FD_ISSET(c.first, &readset)) {
			sockets.push_back({ c.second.first, SocketListener::Read });
		}
		if (FD_ISSET(c.first, &writeset)) {
			sockets.push_back({ c.second.first, SocketListener::Write });
		}
	}

	return sockets;
}

/* --------------------------------------------------------
 * Poll implementation
 * -------------------------------------------------------- */

#if defined(SOCKET_HAVE_POLL)

#if defined(_WIN32)
#  define poll WSAPoll
#endif

short Poll::topoll(int direction) const noexcept
{
	short result(0);

	if (direction & SocketListener::Read) {
		result |= POLLIN;
	}
	if (direction & SocketListener::Write) {
		result |= POLLOUT;
	}

	return result;
}

int Poll::todirection(short event) const noexcept
{
	int direction{};

	/*
	 * Poll implementations mark the socket differently regarding
	 * the disconnection of a socket.
	 *
	 * At least, even if POLLHUP or POLLIN is set, recv() always
	 * return 0 so we mark the socket as readable.
	 */
	if ((event & POLLIN) || (event & POLLHUP)) {
		direction |= SocketListener::Read;
	}
	if (event & POLLOUT) {
		direction |= SocketListener::Write;
	}

	return direction;
}

void Poll::set(Socket &s, int direction) override
{
	auto it = std::find_if(m_fds.begin(), m_fds.end(), [&] (const auto &pfd) { return pfd.fd == s.handle(); });

	// If found, add the new direction, otherwise add a new socket
	if (it != m_fds.end()) {
		it->events |= topoll(direction);
	} else {
		m_lookup.insert({s.handle(), s});
		m_fds.push_back({ s.handle(), topoll(direction), 0 });
	}
}

void Poll::unset(Socket &s, int direction) override
{
	for (auto i = m_fds.begin(); i != m_fds.end();) {
		if (i->fd == s.handle()) {
			i->events &= ~(topoll(direction));

			if (i->events == 0) {
				m_lookup.erase(i->fd);
				i = m_fds.erase(i);
			} else {
				++i;
			}
		} else {
			++i;
		}
	}
}

void Poll::remove(Socket &s) override
{
	auto it = std::find_if(m_fds.begin(), m_fds.end(), [&] (const auto &pfd) { return pfd.fd == s.handle(); });

	if (it != m_fds.end()) {
		m_fds.erase(it);
		m_lookup.erase(s.handle());
	}
}

void Poll::clear() override
{
	m_fds.clear();
	m_lookup.clear();
}

SocketStatus Poll::wait(int ms) override
{
	auto result = poll(m_fds.data(), m_fds.size(), ms);
	if (result == 0) {
		throw SocketError(SocketError::Timeout, "select", "Timeout while listening");
	}
	if (result < 0) {
		throw SocketError(SocketError::System, "poll");
	}

	for (auto &fd : m_fds) {
		if (fd.revents != 0) {
			return { m_lookup.at(fd.fd), todirection(fd.revents) };
		}
	}

	throw SocketError(SocketError::System, "select", "No socket found");
}

std::vector<SocketStatus> Poll::waitMultiple(int ms) override
{
	auto result = poll(m_fds.data(), m_fds.size(), ms);
	if (result == 0) {
		throw SocketError(SocketError::Timeout, "select", "Timeout while listening");
	}
	if (result < 0) {
		throw SocketError(SocketError::System, "poll");
	}

	std::vector<SocketStatus> sockets;
	for (auto &fd : m_fds) {
		if (fd.revents != 0) {
			sockets.push_back({ m_lookup.at(fd.fd), todirection(fd.revents) });
		}
	}

	return sockets;
}

#endif // !_SOCKET_HAVE_POLL

} // !backend