#pragma once
#include "OS/Sync.h"

namespace storm {
	STORM_PKG(core.lang);

	class ActiveFunctions;


	/**
	 * RAII-type used for pausing all threads in the system. Also collects stack-traces for the
	 * threads, to allow querying and replacing active functions on the stack.
	 *
	 * This class is intentionally only usable from C++. The stored stack traces can, however, be
	 * accessed from Storm as well.
	 */
	class PauseThreads : NoCopy {
	public:
		// Pause all threads, capture call-stack information.
		PauseThreads(Engine &e);

		// Release the threads again.
		~PauseThreads();

		// Get active functions. Does not add a ref to the result.
		inline ActiveFunctions *activeFunctions() const {
			return data;
		}

	private:
		// Semaphore used by all threads to signal that they have started waiting.
		// Note: this needs to be a Semaphore, otherwise we might have stack traces for threads
		// that are no longer active.
		Semaphore signal;

		// Use by this class to pause the other threads. Note: it is a Semaphore, not an os::Sema.
		// This pauses the user-level scheduler fully.
		Semaphore wait;

		// Number of threads.
		size_t threadCount;

		// Contents maintained by us.
		ActiveFunctions *data;

		// Lock to avoid concurrent calls to 'data->addThread'.
		// As with other primitives, a native lock to avoid unintended UThread execution.
		util::Lock dataLock;

		// Main entry-point in each thread.
		void threadMain();

		// Capture a stack trace of the current thread.
		void captureStacks(Bool includeCurrent);
	};


	/**
	 * Data returned by 'find' in 'ActiveFunctions'. Contains offset + number of occurrences for a
	 * single active function.
	 */
	struct ActiveOffset {
		// Offset into the queried active function.
		size_t offset;

		// Number of occurrences in the system.
		size_t count;

		// Create.
		ActiveOffset(size_t offset, size_t count)
			: offset(offset), count(count) {}
	};

	// Output.
	wostream &operator <<(wostream &to, const ActiveOffset &offset);

	/**
	 * Data produced by PauseThreads. Includes information about the call stacks for all threads,
	 * including the one that was paused. This allows replacing return pointers on the stack during
	 * updates, for example.
	 *
	 * This data structure is allocated on the C++ heap. All pointers inside it either refer to
	 * locations on the stack, or objects that are pinned on the stack, so it is fine to keep them
	 * in non-gc memory.
	 *
	 * The data structure is cleared whenever the creating PauseThreads object is destroyed. This
	 * prevents stale pointers to the stack.
	 */
	class ActiveFunctions {
	public:
		// Create.
		ActiveFunctions();

		// Add a reference to the refcount.
		void addRef();

		// Release a reference.
		void release();

		// Find all active instances of a particular function.
		vector<ActiveOffset> find(const void *function) const;

		// Replace occurrences of a particular function. Only replaces the specified offsets in the
		// function. Replacements will *not* be reflected in future calls to "find". Returns the
		// number of instances that was replaced. Both 'function' and 'replace' are expected to be
		// pointers to the start of functions allocated on the GC heap.
		size_t replace(const void *function, size_t fOffset, const void *replace, size_t rOffset) const;

	private:
		// This is the API that PauseThreads is expected to use.
		friend class PauseThreads;

		// Add a set of UThread stack traces from an array.
		void addThread(const vector<::StackTrace> &src);

		// Called when data collection is complete to allow pre-processing for efficient lookup down
		// the line.
		void done();

		// Called to clear contents whenever the PauseThreads object goes out of scope.
		void clear();

	private:
		// All stack frames, in one big array to make it easier to refer to them.
		vector<StackFrame> frames;

		// Starting positions of each UThread inside 'frames'.
		vector<size_t> uthreadStart;

		// Starting positions of each Thread inside 'uthreadStart'.
		vector<size_t> threadStart;

		// The elements in 'frames', sorted based on the address in each frame to make them easy to
		// find later on.
		vector<size_t> sortedFrames;

		// Number of references to this structure. Used to avoid stale references from Storm
		// objects.
		size_t refs;
	};

}
