#include <Source/RuntimeInternal.hpp>
#include "ProcessMap.hpp"

#if defined(AURORA_PLATFORM_WIN32)
	#include "ProcessMap.Win32.hpp"
#endif

#if defined(AURORA_IS_MODERNNT_DERIVED)
	#include "ProcessMap.NT.hpp"
#endif

namespace Aurora::Process
{
	struct ModuleBasePairUtil
	{
		AuUInt operator()(const ModuleBasePair &in) const
		{
			return in.modBase;
		}

		constexpr bool operator()(const ModuleBasePair &lhs, const AuString &rhs) const
		{
			return lhs.module == rhs;
		}

		constexpr bool operator()(const ModuleBasePair &lhs, const AuUInt &rhs) const
		{
			return lhs.modBase == rhs;
		}

		constexpr bool operator()(const ModuleBasePair &lhs, const ModuleBasePair &rhs) const
		{
			return rhs.modBase ? lhs.modBase == rhs.modBase : (rhs.module.size() ? rhs.module == rhs.module : false);
		}
	};
	
	struct ModuleLookup : Segment
	{
		ModuleLookup()
		{}

		ModuleLookup(const Segment &s) : Segment(s)
		{}
	};

	static AuBST<AuUInt, ModuleLookup> gModulePtrMap;

	static const auto kMinPageAlignment = 4096;
	static const auto kPageBufferPad = 20;

	static AuThreadPrimitives::MutexUnique_t gMutexUnique;
	static AuHashMapEx<ModuleBasePair, AuSPtr<PublicModule>, ModuleBasePairUtil> gModuleMap;

	static AuUInt ToLowestPageAlignment(AuUInt in)
	{
		return in & ~(kMinPageAlignment - 1);
	}

	static AuSPtr<PublicModule> GetModuleFromSegmentCache(AuUInt pointer)
	{
		AU_LOCK_GUARD(gMutexUnique);
		auto itr = gModulePtrMap.find(pointer);
		if (itr == gModulePtrMap.end()) return {};
		return itr->second.moduleMeta.lock();
	}

	static AuOptional<Segment>  GetSegmentCache(AuUInt pointer)
	{
		AU_LOCK_GUARD(gMutexUnique);
		auto itr = gModulePtrMap.find(pointer);
		if (itr == gModulePtrMap.end()) return {};
		return itr->second;
	}

	bool IsInModuleCache(const ModuleBasePair &pair)
	{
		AU_LOCK_GUARD(gMutexUnique);
		return gModuleMap.find(pair) != gModuleMap.end();
	}

	void InsertModuleCache(const ModuleBasePair &pair, const AuSPtr<PublicModule> &mod)
	{
		for (auto &segment : mod->segments)
		{
			AU_LOCK_GUARD(gMutexUnique);
			segment.moduleMeta = mod;

			for (AuUInt i = segment.baseVa; i < segment.baseVa + segment.size; i += (kMinPageAlignment * kPageBufferPad))
			{
				ModuleLookup a(segment);
				gModulePtrMap[i] = a;
			}
		}

		{
			AU_LOCK_GUARD(gMutexUnique);
			gModuleMap[pair] = mod;
		}
	}

	void RemoveModuleCache(const ModuleBasePair &eitherOr)
	{
		AU_LOCK_GUARD(gMutexUnique);
		auto itr = gModuleMap.find(eitherOr);
		if (itr == gModuleMap.end()) return;
		auto mod = itr->second;

		for (const auto &segment : mod->segments)
		{
			for (AuUInt i = segment.baseVa; i < segment.baseVa + segment.size; i += (kMinPageAlignment * kPageBufferPad))
			{
				auto itr = gModulePtrMap.find(i);
				if (itr != gModulePtrMap.end())
				{
					gModulePtrMap.erase(itr);
				}
			}
		}

		gModuleMap.erase(itr);
	}

	PublicModule GetFromModuleCache(AuUInt handle)
	{
		AU_LOCK_GUARD(gMutexUnique);
		auto itr = gModuleMap.find({"", handle});
		if (itr == gModuleMap.end()) return {};
		return *itr->second;
	}

	static AuOptional<Segment>  FindInCache(AuUInt pointer)
	{
		auto curPtr = ToLowestPageAlignment(pointer);
		auto temp = GetSegmentCache(curPtr);
		if (temp.has_value()) return temp;

		for (int i = 0; i < kPageBufferPad + 1; i++)
		{
			curPtr -= kMinPageAlignment;
			temp = GetSegmentCache(curPtr); // TODO: i dont want to start from the top of the tree, thats stupid
			if (temp.has_value()) return temp;
		}

		return AuOptional<Segment>{};
	}

	void InitProcessMap()
	{
		gMutexUnique = AuThreadPrimitives::MutexUnique();

	#if defined(AURORA_IS_MODERNNT_DERIVED)
		InitProcessMapNt();
	#endif

	#if defined(AURORA_PLATFORM_WIN32)
		MakeToolHelp32Snapshot();
	#endif
	}

	void DeinitProcessMap()
	{
		gMutexUnique.reset();
		DeinitProcessMapNt();
	}

	AUKN_SYM AuOptional<Segment>  GetSegment(AuUInt pointer)
	{
		try 
		{
			auto ceg = FindInCache(pointer);
			if (ceg.has_value())
			{
				return ceg;
			}

		#if defined(AURORA_IS_MODERNNT_DERIVED)
			if (MakeAwarePtr(pointer))
			{
				return FindInCache(pointer);
			}
			else
			{
				return LookupArbitrarySegment(pointer);
			}
		#endif

			return {};
		}
		catch (...)
		{

		}
	}

	AUKN_SYM PublicModule DumpExecutableRoot()
	{
		try
		{
			return GetExecutableRoot();
		}
		catch (...)
		{

		}
	}

	AUKN_SYM Segments DumpExecutableAll()
	{
		try
		{
			Segments ret;
			for (const auto &[meta, ptr] : gModuleMap)
			{
				ret.insert(ret.end(), ptr->segments.begin(), ptr->segments.end());
			}
			return ret;
		}
		catch (...)
		{

		}
	}
}