/*
    This file is part of Mitsuba, a physically based rendering system.

    Copyright (c) 2007-2014 by Wenzel Jakob and others.

    Mitsuba is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.

    Mitsuba is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#include <mitsuba/core/lock.h>
#include <mitsuba/core/fresolver.h>
#include <mitsuba/core/atomic.h>

#if defined(MTS_OPENMP)
# include <omp.h>
#endif

#include <boost/thread/thread.hpp>

// Required for native thread functions
#if defined(__LINUX__)
# include <sys/prctl.h>
#elif defined(__OSX__)
# include <pthread.h>
#elif defined(__WINDOWS__)
# include <windows.h>
#endif

MTS_NAMESPACE_BEGIN

#if defined(_MSC_VER)
namespace {
// Helper function to set a native thread name. MSDN:
//   http://msdn.microsoft.com/en-us/library/xcb2z8hs.aspx

const DWORD MS_VC_EXCEPTION = 0x406D1388;

#pragma pack(push, 8)
struct THREADNAME_INFO {
	DWORD dwType;     // Must be 0x1000.
	LPCSTR szName;    // Pointer to name (in user addr space).
	DWORD dwThreadID; // Thread ID (-1=caller thread).
	DWORD dwFlags;    // Reserved for future use, must be zero.
};
#pragma pack(pop)

void SetThreadName(const char* threadName, DWORD dwThreadID = -1) {
	THREADNAME_INFO info;
	info.dwType     = 0x1000;
	info.szName     = threadName;
	info.dwThreadID = dwThreadID;
	info.dwFlags    = 0;

	__try {
		RaiseException( MS_VC_EXCEPTION, 0, sizeof(info)/sizeof(ULONG_PTR),
			(ULONG_PTR*)&info );
	} __except(EXCEPTION_EXECUTE_HANDLER) { }
}


} // namespace
#endif // _MSC_VER

#if defined(__LINUX__) || defined(__OSX__)
static pthread_key_t __thread_id;
#elif defined(__WINDOWS__)
__declspec(thread) int __thread_id;
#endif
static int __thread_id_ctr = -1;

/**
 * Internal Thread members
 */
struct Thread::ThreadPrivate {
	ref<Thread> parent;
	ref<Logger> logger;
	ref<FileResolver> fresolver;
	boost::mutex joinMutex;
	std::string name;
	bool running, joined;
	Thread::EThreadPriority priority;
	int coreAffinity;
	static ThreadLocal<Thread> *self;
	bool critical;
	boost::thread thread;
	#if defined(__LINUX__) || defined(__OSX__)
		pthread_t native_handle;
	#endif

	ThreadPrivate(const std::string & name_) :
		name(name_), running(false), joined(false),
		priority(Thread::ENormalPriority), coreAffinity(-1),
	    critical(false) { }
};

static std::vector<bool (*)(void)> __crashHandlers;
static std::vector<Thread *> __unmanagedThreads;
static boost::mutex __unmanagedMutex;

/**
 * Dummy class to associate a thread identity with the main thread
 */
class MainThread : public Thread {
public:
	MainThread() : Thread("main") { }

	virtual void run() {
		Log(EError, "The main thread is already running!");
	}

	MTS_DECLARE_CLASS()
protected:
	virtual ~MainThread() { }
};

class UnmanagedThread : public Thread {
public:
	UnmanagedThread(const std::string &name)
		: Thread(name) { }

	virtual void run() {
		Log(EError, "The unmanaged thread is already running!");
	}

	MTS_DECLARE_CLASS()
protected:
	virtual ~UnmanagedThread() { }
};


ThreadLocal<Thread> *Thread::ThreadPrivate::self = NULL;

Thread::Thread(const std::string &name)
 : d(new ThreadPrivate(name)) { }

Thread::EThreadPriority Thread::getPriority() const {
	return d->priority;
}

void Thread::setCritical(bool critical) {
	d->critical = critical;
}

bool Thread::getCritical() const {
	return d->critical;
}

int Thread::getID() {
#if defined(__WINDOWS__)
	return __thread_id;
#elif defined(__OSX__) || defined(__LINUX__)
	/* pthread_self() doesn't provide nice increasing IDs, and syscall(SYS_gettid)
	   causes a context switch. Thus, this function uses a thread-local variable
	   to provide a nice linearly increasing sequence of thread IDs */
	return static_cast<int>(reinterpret_cast<intptr_t>(pthread_getspecific(__thread_id)));
#endif
}

const std::string& Thread::getName() const {
	return d->name;
}

void Thread::setName(const std::string &name) {
	d->name = name;
}

Thread* Thread::getParent() {
	return d->parent;
}

const Thread* Thread::getParent() const {
	return d->parent.get();
}

void Thread::setLogger(Logger *logger) {
	d->logger = logger;
}

Logger* Thread::getLogger() {
	return d->logger;
}

void Thread::setFileResolver(FileResolver *fresolver) {
	d->fresolver = fresolver;
}

FileResolver* Thread::getFileResolver() {
	return d->fresolver;
}

Thread* Thread::getThread() {
	return ThreadPrivate::self->get();
}

bool Thread::isRunning() const {
	return d->running;
}

void Thread::start() {
	if (d->running)
		Log(EError, "Thread is already running!");
	if (!d->self)
		Log(EError, "Threading has not been initialized!");

	Log(EDebug, "Spawning thread \"%s\"", d->name.c_str());

	d->parent = Thread::getThread();

	/* Inherit the parent thread's logger if none was set */
	if (!d->logger)
		d->logger = d->parent->getLogger();

	/* Inherit the parent thread's file resolver if none was set */
	if (!d->fresolver)
		d->fresolver = d->parent->getFileResolver();

	d->running = true;
	d->joined = false;

	incRef();
	try {
		d->thread = boost::thread(&Thread::dispatch, this);
	} catch (boost::thread_resource_error &ex) {
		Log(EError, "Could not create thread!");
		throw ex;
	}
}

bool Thread::setPriority(EThreadPriority priority) {
	d->priority = priority;
	if (!d->running)
		return true;

#if defined(__LINUX__) || defined(__OSX__)
	Float factor;
	switch (priority) {
		case EIdlePriority: factor = 0.0f; break;
		case ELowestPriority: factor = 0.2f; break;
		case ELowPriority: factor = 0.4f; break;
		case EHighPriority: factor = 0.6f; break;
		case EHighestPriority: factor = 0.8f; break;
		case ERealtimePriority: factor = 1.0f; break;
		default: factor = 0.0f; break;
	}

	const pthread_t threadID = d->native_handle;
	struct sched_param param;
	int policy;
	int retval = pthread_getschedparam(threadID, &policy, &param);
	if (retval) {
		Log(EWarn, "pthread_getschedparam(): %s!", strerror(retval));
		return false;
	}

	int min = sched_get_priority_min(policy);
	int max = sched_get_priority_max(policy);

	if (min == max) {
		Log(EWarn, "Could not adjust the thread priority -- valid range is zero!");
		return false;
	}
	param.sched_priority = (int) (min + (max-min)*factor);

	retval = pthread_setschedparam(threadID, policy, &param);
	if (retval) {
		Log(EWarn, "Could not adjust the thread priority to %i: %s!",
			param.sched_priority, strerror(retval));
		return false;
	}
#elif defined(__WINDOWS__)
	int win32Priority;
	switch (priority) {
		case EIdlePriority:
			win32Priority = THREAD_PRIORITY_IDLE;
			break;
		case ELowestPriority:
			win32Priority = THREAD_PRIORITY_LOWEST;
			break;
		case ELowPriority:
			win32Priority = THREAD_PRIORITY_BELOW_NORMAL;
			break;
		case EHighPriority:
			win32Priority = THREAD_PRIORITY_ABOVE_NORMAL;
			break;
		case EHighestPriority:
			win32Priority = THREAD_PRIORITY_HIGHEST;
			break;
		case ERealtimePriority:
			win32Priority = THREAD_PRIORITY_TIME_CRITICAL;
			break;
		default:
			win32Priority = THREAD_PRIORITY_NORMAL;
			break;
	}

	// If the function succeeds, the return value is nonzero
	const HANDLE handle = d->thread.native_handle();
	if (SetThreadPriority(handle, win32Priority) == 0) {
		Log(EWarn, "Could not adjust the thread priority to %i: %s!",
			win32Priority, lastErrorText().c_str());
		return false;
	}
#else
	Log(EWarn, "Thread priority not supported by boost::thread version yet");
#endif
	return true;
}

void Thread::setCoreAffinity(int coreID) {
	d->coreAffinity = coreID;
	if (!d->running)
		return;

#if defined(__OSX__)
	/* CPU affinity not supported on OSX */
#elif defined(__LINUX__)
	int nCores = sysconf(_SC_NPROCESSORS_CONF),
	    nLogicalCores = nCores;

	size_t size = 0;
	cpu_set_t *cpuset = NULL;
	int retval = 0;

	/* The kernel may expect a larger cpu_set_t than would
	   be warranted by the physical core count. Keep querying
	   with increasingly larger buffers if the
	   pthread_getaffinity_np operation fails */
	for (int i = 0; i<6; ++i) {
		size = CPU_ALLOC_SIZE(nLogicalCores);
		cpuset = CPU_ALLOC(nLogicalCores);
		if (!cpuset) {
			Log(EWarn, "Thread::setCoreAffinity(): could not allocate cpu_set_t");
			return;
		}

		CPU_ZERO_S(size, cpuset);

		int retval = pthread_getaffinity_np(d->native_handle, size, cpuset);
		if (retval == 0)
			break;

		/* Something went wrong -- release memory */
		CPU_FREE(cpuset);

		if (retval == EINVAL) {
			/* Retry with a larger cpuset */
			nLogicalCores *= 2;
		} else {
			break;
		}
	}

	if (retval) {
		Log(EWarn, "Thread::setCoreAffinity(): pthread_getaffinity_np(): could "
			"not read thread affinity map: %s", strerror(retval));
		return;
	}

	int actualCoreID = -1, available = 0;
	for (int i=0; i<nLogicalCores; ++i) {
		if (!CPU_ISSET_S(i, size, cpuset))
			continue;
		if (available++ == coreID) {
			actualCoreID = i;
			break;
		}
	}

	if (actualCoreID == -1) {
		Log(EWarn, "Thread::setCoreAffinity(): out of bounds: %i/%i cores available, requested #%i!",
			available, nCores, coreID);
		CPU_FREE(cpuset);
		return;
	}

	CPU_ZERO_S(size, cpuset);
	CPU_SET_S(actualCoreID, size, cpuset);

	retval = pthread_setaffinity_np(d->native_handle, size, cpuset);
	if (retval) {
		Log(EWarn, "Thread::setCoreAffinity(): pthread_setaffinity_np: failed: %s", strerror(retval));
		CPU_FREE(cpuset);
		return;
	}

	CPU_FREE(cpuset);
#elif defined(__WINDOWS__)
	int nCores = getCoreCount();
	const HANDLE handle = d->thread.native_handle();

	DWORD_PTR mask;

	if (coreID != -1 && coreID < nCores)
		mask = (DWORD_PTR) 1 << coreID;
	else
		mask = (1 << nCores) - 1;

	if (!SetThreadAffinityMask(handle, mask))
		Log(EWarn, "Thread::setCoreAffinity(): SetThreadAffinityMask : failed");
#endif
}

int Thread::getCoreAffinity() const {
	return d->coreAffinity;
}

void Thread::dispatch(Thread *thread) {
	detail::initializeLocalTLS();

	#if 0
		#if defined(__LINUX__)
			pthread_setspecific(__thread_id, reinterpret_cast<void *>(syscall(SYS_gettid)));
		#elif defined(__OSX__)
			pthread_setspecific(__thread_id, reinterpret_cast<void *>(pthread_mach_thread_np(pthread_self())));
		#endif
	#endif

	int id = atomicAdd(&__thread_id_ctr, 1);
	#if defined(__LINUX__) || defined(__OSX__)
		pthread_setspecific(__thread_id, reinterpret_cast<void *>(id));
		thread->d->native_handle = pthread_self();
	#elif defined(__WINDOWS__)
		__thread_id = id;
	#endif

	Thread::ThreadPrivate::self->set(thread);

	if (thread->getPriority() != ENormalPriority)
		thread->setPriority(thread->getPriority());

	if (!thread->getName().empty()) {
		const std::string threadName = "Mitsuba: " + thread->getName();
#if defined(__LINUX__)
		// Disabled for now, since it is not yet widely available in glibc
		// pthread_setname_np(pthread_self(), threadName.c_str());

		prctl(PR_SET_NAME, threadName.c_str());
#elif defined(__OSX__)
		pthread_setname_np(threadName.c_str());
#elif defined(__WINDOWS__)
		SetThreadName(threadName.c_str());
#endif
	}

	if (thread->getCoreAffinity() != -1)
		thread->setCoreAffinity(thread->getCoreAffinity());

	try {
		thread->run();
	} catch (std::exception &e) {
		ELogLevel warnLogLevel = thread->getLogger()->getErrorLevel() == EError
			? EWarn : EInfo;
		Log(warnLogLevel, "Fatal error: uncaught exception: \"%s\"", e.what());
		if (thread->d->critical)
			_exit(-1);
	}

	thread->exit();
}

void Thread::join() {
	/* Only one call to join() at a time */
	boost::lock_guard<boost::mutex> guard(d->joinMutex);
	if (d->joined)
		return;
	try {
		d->thread.join();
	} catch (boost::thread_interrupted &ex) {
		Log(EError, "Thread::join() - the thread was interrupted");
		throw ex;
	}
	d->joined = true;
}

void Thread::detach() {
	d->thread.detach();
}

void Thread::sleep(unsigned int ms) {
	try {
		boost::this_thread::sleep(boost::posix_time::milliseconds(ms));
	} catch (boost::thread_interrupted &ex) {
		Log(EError, "Thread::sleep(ms) - interrupted!");
		throw ex;
	}
}

void Thread::yield() {
	boost::this_thread::yield();
}

void Thread::exit() {
	Log(EDebug, "Thread \"%s\" has finished", d->name.c_str());
	d->running = false;
	Assert(ThreadPrivate::self->get() == this);
	detail::destroyLocalTLS();
	decRef();
}

std::string Thread::toString() const {
	std::ostringstream oss;
	oss << "Thread[" << endl
		<< "  name = \"" << d->name << "\"," << endl
		<< "  running = " << d->running << "," << endl
		<< "  joined = " << d->joined << "," << endl
		<< "  priority = " << d->priority << "," << endl
		<< "  critical = " << d->critical << endl
		<< "]";
	return oss.str();
}

#if defined(MTS_OPENMP) && defined(__OSX__)
static int __omp_threadCount = 0;
static pthread_key_t __omp_key;
static bool __omp_key_created;

int mts_omp_get_max_threads() {
	/* This function exists to sidestep an annoying
	   implementation bug that causes crashes on OSX */
	return __omp_threadCount;
}

int mts_omp_get_thread_num() {
	return reinterpret_cast<int>(pthread_getspecific(__omp_key));
}
#endif

void Thread::staticInitialization() {
	#if defined(__OSX__)
		__mts_autorelease_init();
		#if defined(MTS_OPENMP)
			__omp_threadCount = omp_get_max_threads();
		#endif
	#endif
	#if defined(__LINUX__) || defined(__OSX__)
		pthread_key_create(&__thread_id, NULL);
	#endif
	detail::initializeGlobalTLS();
	detail::initializeLocalTLS();

	ThreadPrivate::self = new ThreadLocal<Thread>();
	Thread *mainThread = new MainThread();
	mainThread->d->running = true;
	mainThread->d->joined = false;
	mainThread->d->fresolver = new FileResolver();
	ThreadPrivate::self->set(mainThread);
}

Thread *Thread::registerUnmanagedThread(const std::string &name) {
	Thread *thread = NULL;
	try {
		thread = getThread();
	} catch (...) { }

	if (!thread) {
		detail::initializeLocalTLS();
		thread = new UnmanagedThread(name);
		thread->d->running = false;
		thread->d->joined = false;
		thread->incRef();
		ThreadPrivate::self->set(thread);

		boost::lock_guard<boost::mutex> guard(__unmanagedMutex);
		__unmanagedThreads.push_back((UnmanagedThread *) thread);
	}
	return thread;
}

void Thread::registerCrashHandler(bool (*handler)(void)) {
	__crashHandlers.push_back(handler);
}

void Thread::staticShutdown() {
	for (size_t i=0; i<__unmanagedThreads.size(); ++i)
		__unmanagedThreads[i]->decRef();
	__unmanagedThreads.clear();
	getThread()->d->running = false;
	detail::destroyLocalTLS();
	delete ThreadPrivate::self;
	ThreadPrivate::self = NULL;
	detail::destroyGlobalTLS();
#if defined(__LINUX__) || defined(__OSX__)
	pthread_key_delete(__thread_id);
#endif
#if defined(__OSX__)
	#if defined(MTS_OPENMP)
		if (__omp_key_created)
			pthread_key_delete(__omp_key);
	#endif
	__mts_autorelease_shutdown();
#endif
}

void Thread::initializeOpenMP(size_t threadCount) {
#if defined(MTS_OPENMP)
	ref<Logger> logger = Thread::getThread()->getLogger();
	ref<FileResolver> fResolver = Thread::getThread()->getFileResolver();

	#if defined(__OSX__)
		if (!__omp_key_created) {
			pthread_key_create(&__omp_key, NULL);
			__omp_key_created = true;
		}
		__omp_threadCount = threadCount;
	#endif

	if (omp_get_dynamic())
		omp_set_dynamic(0);

	omp_set_num_threads((int) threadCount);

	int counter = 0;

	#pragma omp parallel
	{
		#if defined(__OSX__)
			if (!pthread_getspecific(__omp_key))
				pthread_setspecific(__omp_key, reinterpret_cast<void *>(counter));
		#endif
		detail::initializeLocalTLS();
		Thread *thread = Thread::getThread();
		if (!thread) {
			#pragma omp critical
			{
				thread = new UnmanagedThread(
					formatString("omp%i", counter));
				counter++;
			}
			const std::string threadName = "Mitsuba: " + thread->getName();

			#if defined(__LINUX__)
				prctl(PR_SET_NAME, threadName.c_str());
			#elif defined(__OSX__)
				pthread_setname_np(threadName.c_str());
			#elif defined(__WINDOWS__)
				SetThreadName(threadName.c_str());
			#endif

			int id = atomicAdd(&__thread_id_ctr, 1);
			#if defined(__LINUX__) || defined(__OSX__)
				pthread_setspecific(__thread_id, reinterpret_cast<void *>(id));
			#elif defined(__WINDOWS__)
				__thread_id = id;
			#endif

			thread->d->running = false;
			thread->d->joined = false;
			thread->d->fresolver = fResolver;
			thread->d->logger = logger;
			thread->incRef();
			ThreadPrivate::self->set(thread);

			#pragma omp critical
				__unmanagedThreads.push_back((UnmanagedThread *) thread);
		}
	}
#else
	if (Thread::getThread()->getLogger() != NULL)
		SLog(EWarn, "Mitsuba was compiled without OpenMP support.");
#endif
}

Thread::~Thread() {
	if (d->running)
		Log(EWarn, "Destructor called while thread '%s' was still running", d->name.c_str());
}

MTS_IMPLEMENT_CLASS(Thread, true, Object)
MTS_IMPLEMENT_CLASS(MainThread, false, Thread)
MTS_IMPLEMENT_CLASS(UnmanagedThread, false, Thread)
MTS_NAMESPACE_END
