#include "stdafx.h"
#include "IOHandle.h"
#include "IORequest.h"

#ifdef LINUX_IO_URING
#include <sys/eventfd.h>
#endif

namespace os {

#if defined(WINDOWS)

	IOHandle::IOHandle() : handle(NULL), pending(0) {}

	IOHandle::IOHandle(HANDLE h) : handle(h), pending(0) {}

	HANDLE IOHandle::v() const {
		/**
		 * NOTE: Due to a bug? in windows 7, an IO Completion Port is signaled at all times when it
		 * is not associated with any file handles. Because of this, we keep track of the current
		 * number of pending IO requests and pretend that we do not have an IO Completion Port if
		 * there are no outstanding requests for this thread.
		 */
		if (atomicRead(pending) > 0)
			return handle;
		else
			return NULL;
	}

	void IOHandle::add(Handle h, const ThreadData *id) {
		HANDLE r = CreateIoCompletionPort(h.v(), handle, (ULONG_PTR)id, 1);
		if (r == NULL) {
			// This fails if the handle does not have the OVERLAPPED flag set.
			// PLN(L"ERROR: " << GetLastError());
		} else {
			handle = r;
		}
	}

	void IOHandle::remove(Handle h, const ThreadData *id) {
		// Not needed.
	}

	void IOHandle::notifyAll(const ThreadData *id) const {
		if (!handle)
			return;

		while (true) {
			DWORD bytes = 0;
			ULONG_PTR key = 0;
			OVERLAPPED *request = NULL;
			BOOL ok = GetQueuedCompletionStatus(handle, &bytes, &key, &request, 0);
			int error = GetLastError();

			if (request) {
				// PLN(L"Got status: " << bytes << L", " << key << L", " << request << L", " << ok);
				if ((ULONG_PTR)id == key) {
					IORequest *r = (IORequest *)request;

					if (ok)
						r->complete(nat(bytes));
					else
						r->failed(nat(bytes), error);
				}
			} else {
				// Nothing was dequeued, abort.
				break;
			}
		}
	}

	void IOHandle::close() {
		if (handle)
			CloseHandle(handle);
		handle = NULL;
	}

	void IOHandle::attach() {
		atomicIncrement(pending);
	}

	void IOHandle::detach() {
		assert(atomicRead(pending) > 0);
		atomicDecrement(pending);
	}

#elif defined(POSIX)

	IOHandle::IOHandle() :
		lock(),
#ifdef LINUX_IO_URING
		linuxIO(), eventfd(-1),
#endif
		handles() {

		if (!uringInit())
			pollInit();
	}

	IOHandle::~IOHandle() {
		uringDestroy();
		// Nothing special for the poll implementation.
	}

	IOHandle::Desc IOHandle::desc() {
		util::Lock::L z(lock);
		Desc d = { handles.data(), handles.capacity() + handles.extraData() };
		return d;
	}

#ifdef LINUX_IO_URING

	/**
	 * io_uring implementation
	 */

	bool IOHandle::uringInit() {
		if (!linuxIO.available())
			return false;

		// Create an eventfd and set it up to receive notifications.

		// Note: We *could* utilize the fact that IOCondition already contains an eventfd and use
		// that directly instead. We would, however, need a fallback to using our own eventfd to be
		// compatible with the Gtk integration in the Gui library.

		eventfd = ::eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK);
		if (eventfd < 0) {
			perror("Failed to create an eventfd");
			assert(false);
		}

		linuxIO.attachEventfd(eventfd);

		// Add the eventfd to the IO map. This is so that we can return it from desc().
		handles.put(eventfd, POLLIN, null);
		return true;
	}

	void IOHandle::uringDestroy() {
		if (eventfd >= 0)
			::close(eventfd);
		eventfd = -1;
	}

	void IOHandle::add(Handle h, const ThreadData *id) {
		if (!linuxIO.available()) {
			pollAdd(h, id);
			return;
		}

		// Nothing needed for io_uring. We don't need to keep track of all attached fds.
	}

	void IOHandle::remove(Handle h, const ThreadData *id) {
		if (!linuxIO.available()) {
			pollRemove(h, id);
			return;
		}

		util::Lock::L z(lock);

		// Remove outstanding requests for handles by asking for cancellation:
		IORequest *empty = null;
		auto begin = activeFDs.lower_bound(std::make_pair(h.v(), empty));
		auto end = activeFDs.lower_bound(std::make_pair(h.v() + 1, empty));

		for (; begin != end; begin++) {
			// Note: 'cancel' does not modify 'activeFDs'.
			cancel(begin->second);
		}
	}

	void IOHandle::attach(Handle h, IORequest *request) {
		if (!linuxIO.available()) {
			pollAttach(h, request);
			return;
		}

		util::Lock::L z(lock);
		activeRequests.insert(request);
		activeFDs.insert(std::make_pair(h.v(), request));

		// We set user_data to a pointer to the request. Otherwise we won't know what is what.
		// Idea: We could simply put active requests in an array and have user_data be indices into
		// that array. That would be cheaper, but it would make duplicate notifications more
		// noteable, however.
		request->request.user_data = size_t(request);
		linuxIO.submit(request->request);
	}

	void IOHandle::attachAndRemove(Handle h, IORequest *request) {
		if (!linuxIO.available()) {
			pollAttachAndRemove(h, request);
			return;
		}

		util::Lock::L z(lock);

		// We do these two while holding the lock. The key thing is that 'remove' does not block.
		remove(h, null);
		attach(h, request);
	}

	void IOHandle::detach(Handle h, IORequest *request) {
		if (!linuxIO.available()) {
			pollDetach(h, request);
			return;
		}

		util::Lock::L z(lock);
		activeRequests.erase(request);
		activeFDs.erase(std::make_pair(h.v(), request));
	}

	void IOHandle::cancel(IORequest *request) {
		if (!linuxIO.available()) {
			pollCancel(request);
			return;
		}

		util::Lock::L z(lock);
		if (activeRequests.count(request)) {
			struct io_uring_sqe job;
			zeroMem(job);
			job.opcode = IORING_OP_ASYNC_CANCEL;
			job.addr = size_t(request);
			job.user_data = 0;
			linuxIO.submit(job);
		}
	}

	void IOHandle::notifyAll(const ThreadData *id) {
		if (!linuxIO.available()) {
			pollNotifyAll(id);
			return;
		}

		// Read the eventfd to see if we should check the LinuxIO object. This also resets the event
		// to non-signalling.
		uint64_t fromfd = 0;
		ssize_t readResult = read(eventfd, &fromfd, sizeof(fromfd));
		// If 'read' returned 0, or if 'fromfd' (= the old value in the eventfd) was 0, then
		// the event was not signalled (read returned 0 => -EAGAIN or -EWOULDBLOCK).
		if (readResult == 0 || fromfd == 0)
			return;

		// All good, check the LinuxIO implementation!
		struct io_uring_cqe entry;
		while (linuxIO.get(entry)) {
			IORequest *request = reinterpret_cast<IORequest *>(entry.user_data);
			if (request == null) {
				// We use user_data = 0 for internal requests to cancel. We can just ignore them if
				// they are successul.
				if (entry.res) {
					WARNING(L"Failing internal io_uring request: " << entry.res);
				}
				continue;
			}
			{
				util::Lock::L z(lock);
				// This has two purposes: we both remove the request from 'active' (we received a
				// response, we don't expect another), and check that we did not get a spurious
				// request from somewhere and write to random locations in memory.
				if (activeRequests.erase(request) == 0) {
					// This is only useful for debugging.
					WARNING(L"Ignoring completion to detached request: " << request);
					continue;
				}
			}

			// We want to do this outside of the lock.
			request->onFinish(&entry.res);
		}
	}

	void IOHandle::close() {
		if (eventfd >= 0)
			::close(eventfd);
		eventfd = -1;
		// Nothing to do for poll.
	}

#else

	/**
	 * stub implementation if we should not use io_uring
	 */

	bool IOHandle::uringInit() {
		return false;
	}

	void IOHandle::uringDestroy() {
		// Nothing to do.
	}

	void IOHandle::add(Handle h, const ThreadData *id) {
		pollAdd(h, id);
	}

	void IOHandle::remove(Handle h, const ThreadData *id) {
		pollRemove(h, id);
	}

	void IOHandle::attach(Handle h, IORequest *request) {
		pollAttach(h, request);
	}

	void IOHandle::attachAndRemove(Handle h, IORequest *request) {
		pollAttachAndRemove(h, request);
	}

	void IOHandle::detach(Handle h, IORequest *request) {
		pollDetach(h, request);
	}

	void IOHandle::cancel(IORequest *request) {
		pollCancel(request);
	}

	void IOHandle::notifyAll(const ThreadData *id) {
		pollNotifyAll(id);
	}

	void IOHandle::close() {
		// Nothing to do for poll.
	}

#endif

	/**
	 * poll implementation as a fallback
	 */

	void IOHandle::pollInit() {}

	void IOHandle::pollAdd(Handle h, const ThreadData *id) {
		UNUSED(h); UNUSED(id);
		// Nothing to do for poll.
	}

	void IOHandle::pollRemove(Handle h, const ThreadData *id) {
		UNUSED(id);
		util::Lock::L z(lock);

		// This lets it know that it was closed.
		int code = -ECANCELED;

		// Mark all pending things as 'complete' and remove them.
		for (nat pos = handles.find(h.v()); pos < handles.capacity(); pos = handles.find(h.v())) {
			IORequest *r = handles.valueAt(pos);
			r->onFinish(&code);

			// Remove!
			handles.remove(pos);
		}
	}

	void IOHandle::pollAttach(Handle h, IORequest *request) {
		// Ask the handle to do the request now in order to figure out what kind of request it is,
		// or if we can actually call poll to wait for it.
		IORequest::DoResult r = request->doRequest();
		if (r == IORequest::doDone) {
			request->onFinish(null);
			return;
		}

		util::Lock::L z(lock);
		short events = 0;
		if (r == IORequest::doWaitRead)
			events |= POLLIN;
		if (r == IORequest::doWaitWrite)
			events |= POLLOUT;
		handles.put(h.v(), events, request);
	}

	void IOHandle::pollAttachAndRemove(Handle h, IORequest *request) {
		// Note: This is typically only done for close.
		util::Lock::L z(lock);
		pollRemove(h, null);
		pollAttach(h, request);
	}

	void IOHandle::pollDetach(Handle h, IORequest *request) {
		util::Lock::L z(lock);
		for (nat pos = handles.find(h.v()); pos < handles.capacity(); pos = handles.next(pos)) {
			if (handles.valueAt(pos) == request) {
				handles.remove(pos);
				break;
			}
		}
	}

	void IOHandle::pollNotifyAll(const ThreadData *id) {
		UNUSED(id);
		util::Lock::L z(lock);

		struct pollfd *wait = handles.data();

		size_t extra = handles.extraData();

		// See if we need to ask the OS for new events...
		bool any = false;
		for (size_t i = 0; i < handles.capacity(); i++) {
			struct pollfd &p = wait[i + extra];
			if (p.fd >= 0 && p.events != 0 && p.revents != 0)
				any = true;
		}

		if (!any) {
			// Find new events.
			int r = poll(wait + extra, handles.capacity(), 0);

			// Any use checking for events?
			if (r <= 0)
				return;
		}

		for (size_t i = 0; i < handles.capacity(); i++) {
			struct pollfd &p = wait[i + extra];

			if (p.events && p.revents) {
				p.revents = 0;

				// Something happened!
				if (IORequest *r = handles.valueAt(i)) {
					IORequest::DoResult res = r->doRequest();
					p.events = 0;
					if (res == IORequest::doWaitRead)
						p.events = POLLIN;
					if (res == IORequest::doWaitWrite)
						p.events = POLLOUT;

					// Note: If we tell it that it is done, it might remove itself right here and
					// now (although the delete flag is mostly used for close()). The only downside
					// of that is that we might miss some handles that were ready. But we will catch
					// them next time we get called (poll will still mark them as ready).
					if (res == IORequest::doDone)
						r->onFinish(null);
				} else {
					// If it did not have an io request, stop listening to it.
					p.events = 0;
				}
			}
		}
	}

	void IOHandle::pollCancel(IORequest *request) {
		// Detach the handle so it will not be signaled later.
		pollDetach(Handle(request->request.fd), request);

		// Tell it that the operation was canceled.
		int res = -ECANCELED;
		request->onFinish(&res);
	}


#endif

}
