#pragma once
#include "Object.h"
#include "TObject.h"
#include "CloneEnv.h"
#include "StrBuf.h"
#include "OS/FnCall.h"
#include "Utils/Templates.h"

namespace storm {
	STORM_PKG(core);

	/**
	 * Function pointers for storm and C++.
	 *
	 * A function pointer here may also contain a this-pointer as well, so that member functions
	 * appear like a regular free function without the first this-parameter.
	 *
	 * Function pointers have slightly different calls semantics compared to regular function calls
	 * in Storm; for function pointers, it is determined entirely at runtime if the parameters are
	 * copied or not. This is because we do not have the same amount of information when dealing
	 * with function pointers.
	 *
	 * Function pointers to members of values are not yet supported.
	 *
	 * TODO: Implement equality check.
	 */

	class FnBase;

	/**
	 * Raw function call returned from 'rawCall'. Not actually usable from Storm, but wraps a raw
	 * C++ function pointer for C++.
	 *
	 * Note: This type does not have a proper constructor in C++; it is created by the Storm
	 * overloads of FnBase in assembler.
	 */
	class RawFn {
		STORM_VALUE;
	public:
		// Create a null function that crashes when 'call' is invoked.
		RawFn();

		// Call the raw function.
		void call(FnBase *ptr, void *out, const void *params) const;

	private:
		// Call this thunk.
		UNKNOWN(PTR_GC) const void *fn;
	};


	/**
	 * Target of a function call. Implemented as an abstract class here, so we do not need to expose
	 * Ref and RefSource to everyone.
	 *
	 * These are too low-level to be exposed to Storm.
	 *
	 * Any subclasses may only contain one pointer-sized data member (which may not contain anything
	 * other than a pointer).
	 */
	class FnTarget {
	public:
		// Clone this target to somewhere else.
		virtual void cloneTo(void *to, size_t size) const = 0;

		// Get the pointer we're pointing to.
		virtual const void *ptr() const = 0;

		// Add some kind of label to a StrBuf for the string representation.
		virtual void toS(StrBuf *to, const RootObject *thisPtr) const = 0;

		// Get the underlying constructor (usually the same as 'ptr', but not always).
		virtual const void *ctor(const RootObject *thisPtr) const = 0;
	};

	/**
	 * Target for raw function pointers.
	 */
	class RawFnTarget : public FnTarget {
	public:
		RawFnTarget(const void *ptr);

		virtual void cloneTo(void *to, size_t size) const;
		virtual const void *ptr() const;
		virtual void toS(StrBuf *to, const RootObject *thisPtr) const;
		virtual const void *ctor(const RootObject *thisPtr) const;
	private:
		const void *data;
	};


	/**
	 * Base class for a function pointer.
	 */
	class FnBase : public Object {
		STORM_CLASS;
	public:
		// Create from C++.
		FnBase(const void *fn, const RootObject *thisPtr, Bool member, Thread *thread);

		// Creation mode:
		enum {
			modeFreeFn    = 0x0, // Free function.
			modeMember    = 0x1, // Member function.
			modeCtor      = 0x2, // Constructor of values or classes
			modeActorCtor = 0x3, // Constructor of actors (involves threading)
		};

		// Create with a generic target and all information given.
		FnBase(const FnTarget &target, const RootObject *thisPtr, Nat mode, Thread *thread);

		// Copy.
		FnBase(const FnBase &o);

		// Deep copy.
		virtual void STORM_FN deepCopy(CloneEnv *env);

		// To string.
		virtual void STORM_FN toS(StrBuf *to) const;


		/**
		 * Call this function:
		 */

		// Call our function (low-level, not typesafe).
		// Takes care of delegating to another thread if needed and adds a this pointer if needed.
		// The 'firstInfo' parameter contains the first parameter, tagged to know the contents.
		// If it is zero, then the first parameter was not interesting. If the LSB is 0, then
		// it was a TObject. If the LSB is 1, then it was a Thread.
		// Note: Does *not* copy parameters if needed, but handles the return value properly.
		// Note: if 'params' are statically allocated, make sure it has room for at least one more element!
		template <class R, int C>
		R callRaw(const os::FnCall<R, C> &params, size_t firstInfo, CloneEnv *env) const {
			byte d[sizeof(R)];
			callRawI(d, params, firstInfo, env);
			R *result = (R *)d;
			R copy = *result;
			result->~R();

			if (needsCopy(firstInfo) & 0x2) {
				if (!env)
					env = new (this) CloneEnv();
				cloned(copy, env);
			}
			return copy;
		}

		// Specialization for returning void.
		template <int C>
		void callRaw(const os::FnCall<void, C> &params, size_t firstInfo, CloneEnv *env) const {
			callRawI(null, params, firstInfo, env);
		}

		// Call function with a pointer to the return value. Low-level function used by other generated code.
		// Note: 'env' is a reference that will be initialized if a CloneEnv was created.
		void callRawI(void *output, const os::FnCallRaw &params, size_t firstInfo, CloneEnv *&env) const;

		// Do we need to copy the parameters for this function given the first TObject? Returns 0x1
		// if parameters need to be copied, 0x2 if the result needs to be copied, or 0x3 if both.
		Byte CODECALL needsCopy(size_t firstInfo) const;

		// Get a pointer to the underlying constructor if this FnBase refers to a constructor.
		const void *rawCtor() const;


		/**
		 * Call the function using a generic parameter list in the form of an array. The
		 * implementation of this function is generated by Storm for all templated subclasses.
		 */

		virtual RawFn STORM_FN rawCall();


	private:
		// Type of function pointer.
		Nat ptrMode;

		// This pointer.
		UNKNOWN(PTR_GC) const RootObject *thisPtr;

		// Thread to call on.
		Thread *thread;

		// Storage for target. Two words are enough for now.
		enum { targetSize = 2 };
		size_t target0;
		UNKNOWN(PTR_GC) size_t target1;

		// Get target.
		inline FnTarget *target() const { return (FnTarget *)&target0; }

		// Compute which thread we want to run on.
		Thread *runOn(size_t firstInfo) const;
	};

	// Declare the template to Storm.
	STORM_TEMPLATE(Fn, createFn);


	/**
	 * Categorize the first parameter.
	 * Note: Regular overloads have higher prioerty than templates as long as no implicit conversion has to be made.
	 * Downcasts are OK.
	 */
	template <class T>
	inline size_t categorizeFirst(const T &) { return 0; }
	inline size_t categorizeFirst(const TObject &t) { return (size_t)&t; }
	inline size_t categorizeFirst(const TObject *t) { return (size_t)t; }
	inline size_t categorizeFirst(TObject &t) { return (size_t)&t; }
	inline size_t categorizeFirst(TObject *t) { return (size_t)t; }
	inline size_t categorizeFirst(const Thread &t) { return 0x1 | (size_t)&t; }
	inline size_t categorizeFirst(const Thread *t) { return 0x1 | (size_t)t; }
	inline size_t categorizeFirst(Thread &t) { return 0x1 | (size_t)&t; }
	inline size_t categorizeFirst(Thread *t) { return 0x1 | (size_t)t; }


	// Helper macro to add a cloned object to a os::FnParams object. We need to keep the cloned
	// value alive until the function is actually called, so we can not use a function.
#define FN_CLONE(T, to, obj, env)						\
	typename RemoveConst<T>::Type c_ ## obj = obj;		\
	cloned(c_ ## obj, env)

	/**
	 * C++ implementation. Supports up to two parameters.
	 *
	 * Note: there is an additional restriction from C++. We can not call functions with reference parameters.
	 */
	template <class R, class P1 = void, class P2 = void>
	class Fn : public FnBase {
		STORM_SPECIAL;
	public:
		// Get the Storm type.
		static Type *stormType(Engine &e) {
			return runtime::cppTemplate(e, FnId, 3, StormInfo<R>::id(), StormInfo<P1>::id(), StormInfo<P2>::id());
		}

		// Create.
		Fn(R (CODECALL *ptr)(P1, P2), Thread *thread = null) : FnBase(address(ptr), null, false, thread) {
			runtime::setVTable(this);
		}

		// Note: parameter 'z' is for SFINAE to ensure that P1 is Q*
		template <class Q>
		Fn(R (CODECALL Q::*ptr)(P2), Thread *thread = null, Q *z = P1()) : FnBase(address(ptr), null, true, thread) {
			runtime::setVTable(this);
		}

		template <class Q>
		Fn(R (CODECALL Q::*ptr)(P1, P2), const Q *obj) : FnBase(address(ptr), obj, true, null) {
			runtime::setVTable(this);
		}

		template <class B>
		Fn(R (CODECALL *ptr)(B *, P1, P2), B *obj, Thread *thread = null) : FnBase(address(ptr), obj, false, thread) {
			runtime::setVTable(this);
		}

		// Copy:
		Fn(const Fn<R, P1, P2> &other) : FnBase(other) {
			runtime::setVTable(this);
		}

		// Call the function.
		R call(P1 p1, P2 p2) const {
			size_t firstInfo = categorizeFirst(p1);

			if (needsCopy(firstInfo) & 0x1) {
				CloneEnv *env = new (this) CloneEnv();
				FN_CLONE(P1, params, p1, env);
				FN_CLONE(P2, params, p2, env);
				os::FnCall<R, 2> params = os::fnCall().add(c_p1).add(c_p2);
				return callRaw(params, firstInfo, env);
			} else {
				os::FnCall<R, 2> params = os::fnCall().add(p1).add(p2);
				return callRaw(params, firstInfo, null);
			}
		}
	};

	/**
	 * 1 parameter.
	 */
	template <class R, class P1>
	class Fn<R, P1, void> : public FnBase {
		STORM_SPECIAL;
	public:
		// Get the Storm type.
		static Type *stormType(Engine &e) {
			return runtime::cppTemplate(e, FnId, 2, StormInfo<R>::id(), StormInfo<P1>::id());
		}

		// Create.
		Fn(R (CODECALL *ptr)(P1), Thread *thread = null) : FnBase(address(ptr), null, false, thread) {
			runtime::setVTable(this);
		}

		// Note: parameter 'z' is for SFINAE to ensure that P1 is Q*
		template <class Q>
		Fn(R (CODECALL Q::*ptr)(), Thread *thread = null, Q *z = P1()) : FnBase(address(ptr), null, true, thread) {
			runtime::setVTable(this);
		}

		template <class Q>
		Fn(R (CODECALL Q::*ptr)(P1), const Q *obj) : FnBase(address(ptr), obj, true, null) {
			runtime::setVTable(this);
		}

		template <class B>
		Fn(R (CODECALL *ptr)(B *, P1), B *obj, Thread *thread = null) : FnBase(address(ptr), obj, false, thread) {
			runtime::setVTable(this);
		}

		Fn(const FnTarget &target, RootObject *thisPtr, Nat mode) : FnBase(target, thisPtr, mode, null) {
			runtime::setVTable(this);
		}

		// Copy:
		Fn(const Fn<R, P1> &other) : FnBase(other) {
			runtime::setVTable(this);
		}

		// Call the function.
		R call(P1 p1) const {
			size_t firstInfo = categorizeFirst(p1);

			if (needsCopy(firstInfo) & 0x1) {
				CloneEnv *env = new (this) CloneEnv();
				FN_CLONE(P1, params, p1, env);
				os::FnCall<R, 1> params = os::fnCall().add(c_p1);
				return callRaw(params, firstInfo, env);
			} else {
				os::FnCall<R, 1> params = os::fnCall().add(p1);
				return callRaw(params, firstInfo, null);
			}
		}
	};

	/**
	 * 0 parameters.
	 */
	template <class R>
	class Fn<R, void, void> : public FnBase {
		STORM_SPECIAL;
	public:
		// Get the Storm type.
		static Type *stormType(Engine &e) {
			return runtime::cppTemplate(e, FnId, 1, StormInfo<R>::id());
		}

		// Create.
		Fn(R (CODECALL *ptr)(), Thread *thread = null) : FnBase(address(ptr), null, false, thread) {
			runtime::setVTable(this);
		}

		template <class Q>
		Fn(R (CODECALL Q::*ptr)(), const Q *obj) : FnBase(address(ptr), obj, true, null) {
			runtime::setVTable(this);
		}

		template <class B>
		Fn(R (CODECALL *ptr)(B *), B *obj, Thread *thread = null) : FnBase(address(ptr), obj, false, thread) {
			runtime::setVTable(this);
		}

		// Copy:
		Fn(const Fn<R> &other) : FnBase(other) {
			runtime::setVTable(this);
		}

		// Call the function.
		R call() const {
			// Note: 'callRaw' will create a CloneEnv if it requires one in this case.
			os::FnCall<R, 1> params = os::fnCall();
			return callRaw(params, 0, null);
		}
	};


	/**
	 * Create easily.
	 */

	// Free functions.
	template <class R>
	Fn<R> *fnPtr(Engine &e, R (CODECALL *fn)(), Thread *t = null) {
		return new (e) Fn<R>(fn, t);
	}

	template <class R, class P1>
	Fn<R, P1> *fnPtr(Engine &e, R (CODECALL *fn)(P1), Thread *t = null) {
		return new (e) Fn<R, P1>(fn, t);
	}

	template <class R, class P1, class P2>
	Fn<R, P1, P2> *fnPtr(Engine &e, R (CODECALL *fn)(P1, P2), Thread *t = null) {
		return new (e) Fn<R, P1, P2>(fn, t);
	}


	// Member functions without bound 'this':
	template <class R, class Q>
	Fn<R, Q*> *fnPtr(Engine &e, R (CODECALL Q::*fn)(), Thread *t = null) {
		return new (e) Fn<R, Q*>(fn);
	}

	template <class R, class Q, class P>
	Fn<R, Q*, P> *fnPtr(Engine &e, R (CODECALL Q::*fn)(), Thread *t = null) {
		return new (e) Fn<R, Q*, P>(fn);
	}


	// Member functions with bound 'this':
	template <class R, class Q>
	Fn<R> *fnPtr(Engine &e, R (CODECALL Q::*fn)(), const Q *obj) {
		return new (e) Fn<R>(fn, obj);
	}

	template <class R, class Q, class P1>
	Fn<R, P1> *fnPtr(Engine &e, R (CODECALL Q::*fn)(P1), const Q *obj) {
		return new (e) Fn<R, P1>(fn, obj);
	}

	template <class R, class Q, class P1, class P2>
	Fn<R, P1, P2> *fnPtr(Engine &e, R (CODECALL Q::*fn)(P1, P2), const Q *obj) {
		return new (e) Fn<R, P1, P2>(fn, obj);
	}


	// Free functions, but bind first parameter:
	template <class R, class B>
	Fn<R> *fnBoundPtr(Engine &e, R (CODECALL *fn)(B *), B *bound, Thread *t = null) {
		return new (e) Fn<R>(fn, bound, t);
	}

	template <class R, class B, class P1>
	Fn<R, P1> *fnBoundPtr(Engine &e, R (CODECALL *fn)(B *, P1), B *bound, Thread *t = null) {
		return new (e) Fn<R, P1>(fn, bound, t);
	}

	template <class R, class B, class P1, class P2>
	Fn<R, P1, P2> *fnBoundPtr(Engine &e, R (CODECALL *fn)(B *, P1, P2), B *bound, Thread *t = null) {
		return new (e) Fn<R, P1, P2>(fn, bound, t);
	}

	namespace impl {
		template <class T, class P1>
		class RawCtorTarget : public FnTarget {
		public:
			RawCtorTarget() {}

			virtual void cloneTo(void *to, size_t size) const {
				new (to) RawCtorTarget<T, P1>();
			}
			virtual const void *ptr() const {
				return address(&RawCtorTarget<T, P1>::alloc);
			}
			virtual const void *ctor(const RootObject *thisPtr) const {
				return address(&RawCtorTarget<T, P1>::init);
			}
			virtual void toS(StrBuf *to, const RootObject *thisPtr) const {
				*to << S("C++ ctor");
			}

		private:
			static T *CODECALL alloc(Engine *e, P1 p1) {
				return new (*e) T(p1);
			}
			static void CODECALL init(void *out, P1 p1) {
				new (Place(out)) T(p1);
			}
		};
	}

	// Create pointers to constructors in C++. Only supports 1 parameter at the moment, for
	// serialization. More complex pointers can be created in Storm.
	template <class T, class P1>
	Fn<T *, P1> *ctorPtr(Engine &e, Thread *t = null) {
		return new (e) Fn<T *, P1>(impl::RawCtorTarget<T, P1>(), (RootObject *)&e, FnBase::modeCtor);
	}

#undef FN_CLONE

}