/*
 * Copyright 2005, Ingo Weinhold, bonefish@users.sf.net.
 * Distributed under the terms of the MIT License.
 */

#include <algorithm>

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/select.h>

#include <OS.h>

#include <util/DoublyLinkedList.h>

//	Test cases:
//	
//	1. unblock on close:
//	
//	* unblock, when the same cookie is closed
//	  - read
//	  - write (with, without ECHO)
//	* unblock write operations, when the other tty is closed
//	* unblock slave operations, when the master is closed
//	
//	
//	2. select:
//	
//	* notify read, write, if ready when select()ing
//	* notify read, write, error on close of the other TTY
//    (the select() behavior when closing the same TTY is undefined)
//	* notify read after pending write
//	  - with ECHO
//	  - without ECHO
//	* notify write after pending read and full buffer
//	  - with ECHO
//	  - without ECHO
//	* don't notify when there was a pending read/write
//
//
// TODO: There are no ECHO tests yet, since I don't know how to turn it on.

#define CHK(condition)		assert(condition)
#define RUN_TEST(test)		(test)->Run()
#define CHK_ALIVE(thread)	CHK((thread)->IsAlive())
#define CHK_DEAD(thread)	CHK(!(thread)->IsAlive())

// FDSet
struct FDSet : fd_set {
	FDSet()
	{
		Clear();
	}

	void Clear()
	{
		FD_ZERO(this);
		fCount = 0;
	}

	void Add(int fd)
	{
		if (fd < 0 || fd >= FD_SETSIZE) {
			fprintf(stderr, "FDSet::Add(%d): Invalid FD.\n", fd);
			return;
		}

		FD_SET(fd, this);
		if (fd >= fCount)
			fCount = fd + 1;
	}

	int Count() const
	{
		return fCount;
	}

	void UpdateCount()
	{
		if (fCount > 0) {
			for (int i = fCount - 1; i >= 0; i--) {
				if (FD_ISSET(i, this)) {
					fCount = i + 1;
					return;
				}
			}
			fCount = 0;
		}
	}

	bool operator==(const FDSet &other) const
	{
		if (fCount != other.fCount)
			return false;

		for (int i = 0; i < fCount; i++) {
			if (FD_ISSET(i, this) != FD_ISSET(i, &other))
				return false;
		}

		return true;
	}

	bool operator!=(const FDSet &other) const
	{
		return !(*this == other);
	}

private:
	int	fCount;
};

// SelectSet
class SelectSet {
public:
	SelectSet() {}
	~SelectSet() {}

	void Clear()
	{
		fReadSet.Clear();
		fWriteSet.Clear();
		fErrorSet.Clear();
	}

	void AddReadFD(int fd)
	{
		fReadSet.Add(fd);
	}

	void AddWriteFD(int fd)
	{
		fWriteSet.Add(fd);
	}

	void AddErrorFD(int fd)
	{
		fErrorSet.Add(fd);
	}

	int Select(bigtime_t timeout = -1)
	{
		int count = max(max(fReadSet.Count(), fWriteSet.Count()),
			fErrorSet.Count());
		fd_set *readSet = (fReadSet.Count() > 0 ? &fReadSet : NULL);
		fd_set *writeSet = (fWriteSet.Count() > 0 ? &fWriteSet : NULL);
		fd_set *errorSet = (fErrorSet.Count() > 0 ? &fErrorSet : NULL);

		timeval tv = { timeout / 1000000, timeout % 1000000 };

		int result = select(count, readSet, writeSet, errorSet,
			(timeout >= 0 ? &tv : NULL));
		if (result <= 0) {
			Clear();
		} else {
			fReadSet.UpdateCount();
			fWriteSet.UpdateCount();
			fErrorSet.UpdateCount();
		}

		return result;
	}

	bool operator==(const SelectSet &other) const
	{
		return (fReadSet == other.fReadSet
			&& fWriteSet == other.fWriteSet
			&& fErrorSet == other.fErrorSet);
	}

	bool operator!=(const SelectSet &other) const
	{
		return !(*this == other);
	}

private:
	FDSet	fReadSet;
	FDSet	fWriteSet;
	FDSet	fErrorSet;
};

// Runnable
class Runnable {
public:
	virtual ~Runnable() {}
	virtual int32 Run() = 0;
};

// Caller
template<typename Type>
class Caller : public Runnable {
public:
	Caller(Type *object, int32 (Type::*function)())
		: fObject(object),
		  fFunction(function)
	{
	}

	virtual int32 Run()
	{
		return (fObject->*fFunction)();
	}

private:
	Type	*fObject;
	int32	(Type::*fFunction)();
};

// create_caller
template<typename Type>
static inline Runnable *
create_caller(Type *object, int32 (Type::*function)())
{
	return new Caller<Type>(object, function);
}

#define CALLER(object, function)	create_caller(object, function)

// Thread
struct Thread : public DoublyLinkedListLinkImpl<Thread> {
public:
	Thread(Runnable *runnable, const char *name)
		: fRunnable(runnable)
	{
		fThread = spawn_thread(_Entry, name, B_NORMAL_PRIORITY, this);
		if (fThread < 0) {
			sprintf("Failed to spawn thread: %s\n", strerror(fThread));
			exit(1);
		}
	}

	~Thread()
	{
		Kill();
		delete fRunnable;
	}

	void Resume()
	{
		resume_thread(fThread);
	}

	void WaitFor()
	{
		status_t result;
		wait_for_thread(fThread, &result);
	}

	void Kill()
	{
		kill_thread(fThread);
	}

	bool IsAlive()
	{
		thread_info info;
		return (get_thread_info(fThread, &info) == B_OK);
	}

private:
	static int32 _Entry(void *data)
	{
		return ((Thread*)data)->fRunnable->Run();
	}

	thread_id	fThread;
	Runnable	*fRunnable;
};

// TestCase
class TestCase {
public:
	TestCase() {}
	virtual ~TestCase()
	{
		while (Thread *thread = fThreads.Head()) {
			fThreads.Remove(thread);
			delete thread;
		}
	}

	void Run()
	{
		Test();
		delete this;
	}

protected:
	virtual void Test() = 0;

	Thread *CreateThread(Runnable *runnable, const char *name)
	{
		Thread *thread = new Thread(runnable, name);
		fThreads.Add(thread);
		return thread;
	}

	static void WriteUntilBlock(int fd)
	{
		// set non-blocking I/O
		if (fcntl(fd, F_SETFL, O_NONBLOCK) < 0) {
			fprintf(stderr, "WriteUntilBlock(): Failed to set non-blocking I/O "
				"mode: %s\n", strerror(errno));
			exit(1);
		}

		// write till blocking
		char buffer[1024];
		memset(buffer, 'A', sizeof(buffer));
		ssize_t bytesWritten;
		do {
			bytesWritten = write(fd, buffer, sizeof(buffer));
		} while (bytesWritten > 0 || errno == B_INTERRUPTED);

		if (bytesWritten < 0 && errno != B_WOULD_BLOCK) {
			fprintf(stderr, "WriteUntilBlock(): Writing failed: %s\n",
				strerror(errno));
			exit(1);
		}

		// reset to blocking I/O
		if (fcntl(fd, F_SETFL, 0) < 0) {
			fprintf(stderr, "WriteUntilBlock(): Failed to set blocking I/O "
				"mode: %s\n", strerror(errno));
			exit(1);
		}
	}

	static void ReadDontFail(int fd, int32 size)
	{
		char buffer[1024];

		while (size > 0) {
			ssize_t bytesRead;
			do {
				int32 toRead = sizeof(buffer);
				if (toRead > size)
					toRead = size;

				bytesRead = read(fd, buffer, toRead);
			} while (bytesRead < 0 && errno == B_INTERRUPTED);

			if (bytesRead < 0) {
				fprintf(stderr, "ReadDontFail(): Failed to read: %s\n",
					strerror(errno));
				exit(1);
			}

			size -= bytesRead;
		}
	}

	static void WriteDontFail(int fd, int32 size)
	{
		char buffer[1024];
		memset(buffer, 'A', sizeof(buffer));

		while (size > 0) {
			ssize_t bytesWritten;
			do {
				int32 toWrite = sizeof(buffer);
				if (toWrite > size)
					toWrite = size;

				bytesWritten = write(fd, buffer, toWrite);
			} while (bytesWritten < 0 && errno == B_INTERRUPTED);

			if (bytesWritten < 0) {
				fprintf(stderr, "WriteDontFail(): Failed to write: %s\n",
					strerror(errno));
				exit(1);
			}

			size -= bytesWritten;
		}
	}

	void SetEcho(int fd)
	{
		// TODO: How to set echo mode?
	}

private:
	typedef DoublyLinkedList<Thread> ThreadList;

	ThreadList	fThreads;
};


// open_tty
static int
open_tty(int index, bool master)
{
	if (index < 0 || index >= 16)
		fprintf(stderr, "open_tty(%d, %d): Bad index!\n", index, master);

	char path[32];
	sprintf(path, "/dev/%ct/r%x", (master ? 'p' : 't'), index);
	int fd = open(path, O_RDWR);

	if (fd < 0) {
		fprintf(stderr, "Failed to open tty `%s': %s\n", path, strerror(errno));
		exit(1);
	}

	return fd;
}


static void
close_tty(int &fd)
{
	if (fd >= 0) {
		close(fd);
		fd = -1;
	}
}


// #pragma mark -

// TestUnblockOnCloseRead
class TestUnblockOnCloseRead : public TestCase {
public:
	TestUnblockOnCloseRead(bool master, bool crossOver)
		: fMaster(-1),
		  fSlave(-1),
		  fTestMaster(master),
		  fCrossOver(crossOver)
	{
		printf("TestUnblockOnCloseRead(%d, %d)\n", master, crossOver);
	}


protected:
	virtual ~TestUnblockOnCloseRead()
	{
		close_tty(fMaster);
		close_tty(fSlave);
	}

	virtual void Test()
	{
		fMaster = open_tty(0, true);
		fSlave = open_tty(0, false);

		Thread *thread = CreateThread(
			CALLER(this, &TestUnblockOnCloseRead::_Reader), "reader");
		thread->Resume();

		snooze(100000);
		CHK_ALIVE(thread);

		if (fCrossOver)
			close_tty(fTestMaster ? fSlave : fMaster);
		else
			close_tty(fTestMaster ? fMaster : fSlave);

		snooze(100000);
		CHK_DEAD(thread);
	}

private:
	int32 _Reader()
	{
		char buffer[32];
		ssize_t bytesRead = read((fTestMaster ? fMaster : fSlave), buffer,
			sizeof(buffer));

		CHK((bytesRead < 0));

		return 0;
	}

private:
	int		fMaster;
	int		fSlave;
	bool	fTestMaster;
	bool	fCrossOver;
};

// TestUnblockOnCloseWrite
class TestUnblockOnCloseWrite : public TestCase {
public:
	TestUnblockOnCloseWrite(bool master, bool crossOver, bool echo)
		: fMaster(-1),
		  fSlave(-1),
		  fTestMaster(master),
		  fCrossOver(crossOver),
		  fEcho(echo)
	{
		printf("TestUnblockOnCloseWrite(%d, %d, %d)\n", master, crossOver,
			echo);
	}


protected:
	virtual ~TestUnblockOnCloseWrite()
	{
		close_tty(fMaster);
		close_tty(fSlave);
	}

	virtual void Test()
	{
		fMaster = open_tty(0, true);
		fSlave = open_tty(0, false);

		if (fEcho)
			SetEcho((fTestMaster ? fSlave : fMaster));

		WriteUntilBlock((fTestMaster ? fMaster : fSlave));

		Thread *thread = CreateThread(
			CALLER(this, &TestUnblockOnCloseWrite::_Writer), "writer");
		thread->Resume();

		snooze(100000);
		CHK_ALIVE(thread);

		if (fCrossOver)
			close_tty(fTestMaster ? fSlave : fMaster);
		else
			close_tty(fTestMaster ? fMaster : fSlave);

		snooze(100000);
		CHK_DEAD(thread);
	}

private:
	int32 _Writer()
	{
		char buffer[32];
		memset(buffer, 'A', sizeof(buffer));
		ssize_t bytesWritten = write((fTestMaster ? fMaster : fSlave), buffer,
			sizeof(buffer));

		CHK((bytesWritten < 0));

		return 0;
	}

private:
	int		fMaster;
	int		fSlave;
	bool	fTestMaster;
	bool	fCrossOver;
	bool	fEcho;
};

// TestSelectAlreadyReady
class TestSelectAlreadyReady : public TestCase {
public:
	TestSelectAlreadyReady(bool master, bool write)
		: fMaster(-1),
		  fSlave(-1),
		  fTestMaster(master),
		  fWrite(write)
	{
		printf("TestSelectAlreadyReady(%d, %d)\n", master, write);
	}


protected:
	virtual ~TestSelectAlreadyReady()
	{
		close_tty(fMaster);
		close_tty(fSlave);
	}

	virtual void Test()
	{
		fMaster = open_tty(0, true);
		fSlave = open_tty(0, false);

		if (!fWrite)
			WriteDontFail((fTestMaster ? fSlave : fMaster), 1);

		Thread *thread = CreateThread(
			CALLER(this, &TestSelectAlreadyReady::_Selector), "selector");
		thread->Resume();

		snooze(100000);
		CHK_DEAD(thread);
	}

private:
	int32 _Selector()
	{
		SelectSet selectSet;
		SelectSet compareSet;
		if (fWrite) {
			selectSet.AddWriteFD((fTestMaster ? fMaster : fSlave));
			compareSet.AddWriteFD((fTestMaster ? fMaster : fSlave));
		} else {
			selectSet.AddReadFD((fTestMaster ? fMaster : fSlave));
			compareSet.AddReadFD((fTestMaster ? fMaster : fSlave));
		}

		int result = selectSet.Select();
		CHK(result > 0);
		CHK(selectSet == compareSet);

		return 0;
	}

private:
	int		fMaster;
	int		fSlave;
	bool	fTestMaster;
	bool	fWrite;
};

// TestSelectNotifyOnClose
class TestSelectNotifyOnClose : public TestCase {
public:
	TestSelectNotifyOnClose(bool master)
		: fMaster(-1),
		  fSlave(-1),
		  fTestMaster(master)
	{
		printf("TestSelectNotifyOnClose(%d)\n", master);
	}


protected:
	virtual ~TestSelectNotifyOnClose()
	{
		close_tty(fMaster);
		close_tty(fSlave);
	}

	virtual void Test()
	{
		fMaster = open_tty(0, true);
		fSlave = open_tty(0, false);

		WriteUntilBlock((fTestMaster ? fMaster : fSlave));

		Thread *thread = CreateThread(
			CALLER(this, &TestSelectNotifyOnClose::_Selector), "selector");
		thread->Resume();

		snooze(100000);
		CHK_ALIVE(thread);

		close_tty((fTestMaster ? fSlave : fMaster));

		snooze(100000);
		CHK_DEAD(thread);
	}

private:
	int32 _Selector()
	{
		int fd = (fTestMaster ? fMaster : fSlave);

		SelectSet selectSet;
		selectSet.AddReadFD(fd);
		selectSet.AddWriteFD(fd);
		selectSet.AddErrorFD(fd);

		// In case the slave is closed while we select() on the master, only a
		// `write' event will arrive.
		SelectSet compareSet;
		if (!fTestMaster) {
			compareSet.AddReadFD(fd);
			compareSet.AddErrorFD(fd);
		}
		compareSet.AddWriteFD(fd);

		int result = selectSet.Select();
		CHK(result > 0);
		CHK(selectSet == compareSet);

		return 0;
	}

private:
	int		fMaster;
	int		fSlave;
	bool	fTestMaster;
};

// TestSelectNotifyAfterPending
class TestSelectNotifyAfterPending : public TestCase {
public:
	TestSelectNotifyAfterPending(bool master, bool write, bool unblock)
		: fMaster(-1),
		  fSlave(-1),
		  fTestMaster(master),
		  fWrite(write),
		  fUnblock(unblock)
	{
		printf("TestSelectNotifyAfterPending(%d, %d, %d)\n", master, write,
			unblock);
	}


protected:
	virtual ~TestSelectNotifyAfterPending()
	{
		close_tty(fMaster);
		close_tty(fSlave);
	}

	virtual void Test()
	{
		fMaster = open_tty(0, true);
		fSlave = open_tty(0, false);

		if (fWrite)
			WriteUntilBlock((fTestMaster ? fMaster : fSlave));

		Thread *readWriter = CreateThread(
			CALLER(this, &TestSelectNotifyAfterPending::_ReadWriter),
			"read-writer");
		Thread *selector = CreateThread(
			CALLER(this, &TestSelectNotifyAfterPending::_Selector), "selector");

		readWriter->Resume();
		selector->Resume();

		snooze(100000);
		CHK_ALIVE(readWriter);
		CHK_ALIVE(selector);

		if (fUnblock) {
			// unblock the read-writer and the selector
			if (fWrite)
				ReadDontFail((fTestMaster ? fSlave : fMaster), 2);
			else
				WriteDontFail((fTestMaster ? fSlave : fMaster), 2);

			snooze(100000);
			CHK_DEAD(readWriter);
			CHK_DEAD(selector);

		} else {
			// unblock the read-writer, but not the selector
			if (fWrite)
				ReadDontFail((fTestMaster ? fSlave : fMaster), 1);
			else
				WriteDontFail((fTestMaster ? fSlave : fMaster), 1);

			snooze(100000);
			CHK_DEAD(readWriter);
			CHK_ALIVE(selector);

			close_tty((fTestMaster ? fMaster : fSlave));

			snooze(100000);
			CHK_DEAD(selector);
		}
	}

private:
	int32 _ReadWriter()
	{
		int fd = (fTestMaster ? fMaster : fSlave);
		if (fWrite)
			WriteDontFail(fd, 1);
		else
			ReadDontFail(fd, 1);

		return 0;
	}

	int32 _Selector()
	{
		int fd = (fTestMaster ? fMaster : fSlave);

		SelectSet selectSet;
		SelectSet compareSet;

		if (fWrite) {
			selectSet.AddWriteFD(fd);
			compareSet.AddWriteFD(fd);
		} else {
			selectSet.AddReadFD(fd);
			compareSet.AddReadFD(fd);
		}

		int result = selectSet.Select();
		CHK(result > 0);
		CHK(selectSet == compareSet);

		return 0;
	}

private:
	int		fMaster;
	int		fSlave;
	bool	fTestMaster;
	bool	fWrite;
	bool	fUnblock;
};



// #pragma mark -

int
main()
{
	// unblock tests

	RUN_TEST(new TestUnblockOnCloseRead(true,  false));
	RUN_TEST(new TestUnblockOnCloseRead(false, false));
	RUN_TEST(new TestUnblockOnCloseRead(false, true));

	RUN_TEST(new TestUnblockOnCloseWrite(true,  false, false));
	RUN_TEST(new TestUnblockOnCloseWrite(false, false, false));
	RUN_TEST(new TestUnblockOnCloseWrite(true,  true,  false));
	RUN_TEST(new TestUnblockOnCloseWrite(false, true,  false));
// TODO: How to enable echo mode?
//	RUN_TEST(new TestUnblockOnCloseWrite(true,  false, true));
//	RUN_TEST(new TestUnblockOnCloseWrite(false, false, true));
//	RUN_TEST(new TestUnblockOnCloseWrite(true,  true,  true));
//	RUN_TEST(new TestUnblockOnCloseWrite(false, true,  true));

	// select tests

	RUN_TEST(new TestSelectAlreadyReady(true,  true));
	RUN_TEST(new TestSelectAlreadyReady(false, true));
	RUN_TEST(new TestSelectAlreadyReady(true,  false));
	RUN_TEST(new TestSelectAlreadyReady(false, false));

	RUN_TEST(new TestSelectNotifyOnClose(true));
	RUN_TEST(new TestSelectNotifyOnClose(false));

	RUN_TEST(new TestSelectNotifyAfterPending(false, false, false));
	RUN_TEST(new TestSelectNotifyAfterPending(false, false, true));
	RUN_TEST(new TestSelectNotifyAfterPending(false, true,  false));
	RUN_TEST(new TestSelectNotifyAfterPending(false, true,  true));
	RUN_TEST(new TestSelectNotifyAfterPending(true,  false, false));
	RUN_TEST(new TestSelectNotifyAfterPending(true,  false, true));
	RUN_TEST(new TestSelectNotifyAfterPending(true,  true,  false));
	RUN_TEST(new TestSelectNotifyAfterPending(true,  true,  true));

	return 0;
}