/* -----------------------------------------------------------------------------
	GSFramework
	2023 Emmanuel Julien
----------------------------------------------------------------------------- */

#include "binding_helpers.h"

#include <chrono>
#include <fstream>
#include <iostream>
#include <map>
#include <set>
#include <string>
#include <vector>

struct CallId {
	SQUserPointer caller; // caller funcid, 0 means native call
	SQUserPointer funcid;
};

typedef uint64_t time_ns;

static time_ns get_clock() {
	return std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::high_resolution_clock::now().time_since_epoch()).count();
}

typedef uint32_t CallIdx;

struct Call {
	CallId id;
	time_ns start;
	time_ns total{0};
	uint32_t hit{0}; // number of time this call was performed
	std::vector<CallIdx> child_calls; // [EJ] this is wasteful and inefficient
};

struct FuncInfo {
	std::string name;
	std::string source;
};

struct VMProfile {
	std::map<SQUserPointer, FuncInfo> func_info;

	std::vector<Call> all_calls;
	CallIdx call_count{0};

	std::vector<CallIdx> root_calls;
	std::vector<CallIdx> callstack; // current callstack
};

static std::map<HSQUIRRELVM, VMProfile> vm_profiles;

static CallIdx find_call(std::vector<Call> &all_calls, std::vector<CallIdx> &calls, CallId call_id) {
	for (CallIdx i : calls) {
		const Call &call = all_calls[i];
		if (call.id.caller == call_id.caller && call.id.funcid == call_id.funcid) {
			return i;
		}
	}
	return 0;
}

struct CallProfile {
	uint32_t hit; // number of calls
	time_ns total; // duration of all calls
	time_ns child; // duration of all child calls
};

static CallProfile get_call_profile(const std::vector<Call> &all_calls, const Call &call) {
	CallProfile profile;
	profile.hit = call.hit;
	profile.total = call.total;
	profile.child = 0;

	for (const CallIdx i : call.child_calls) {
		profile.child += all_calls[i].total;
	}

	return profile;
}

static CallProfile get_calls_profile(const std::vector<Call> &all_calls, const std::vector<CallIdx> &idxs) {
	CallProfile profile;
	profile.hit = 0;
	profile.total = 0;
	profile.child = 0;

	for (const auto idx : idxs) {
		const CallProfile call_profile = get_call_profile(all_calls, all_calls[idx]);
		profile.hit += call_profile.hit;
		profile.total += call_profile.total;
		profile.child += call_profile.child;
	}

	return profile;
}

static void native_hook(HSQUIRRELVM vm, SQInteger event_type, const SQChar *sourcename, SQInteger line, const SQChar *funcname) {
	VMProfile &profile = vm_profiles[vm];

	const time_ns time = get_clock();

	if (event_type == 'l') { // line execution
		// TODO reimplement if ever needed
	} else if (event_type == 'c') { // function call
		SQFunctionInfo fi;
		sq_getfunctioninfo(vm, 0, &fi);

		if (profile.func_info.find(fi.funcid) == std::end(profile.func_info)) {
			FuncInfo &info = profile.func_info[fi.funcid];
			info.name = fi.name;
			info.source = fi.source;
		}

		//
		CallId call_id = {nullptr, fi.funcid};

		std::vector<CallIdx> *child_calls = &profile.root_calls;

		if (!profile.callstack.empty()) {
			Call &call = profile.all_calls[profile.callstack.back()];

			call_id.caller = call.id.funcid;
			child_calls = &call.child_calls;
		}

		//
		CallIdx i = find_call(profile.all_calls, *child_calls, call_id);

		if (i == 0) {
			i = ++profile.call_count;
			child_calls->push_back(i);

			profile.all_calls.resize(profile.call_count + 1); // allocate call object
		}

		Call &call = profile.all_calls[i];
		call.id = call_id;
		call.start = time;
		++call.hit;

		profile.callstack.push_back(i);
	} else if (event_type == 'r') { // returning from a function
		if (!profile.callstack.empty()) {
			Call &call = profile.all_calls[profile.callstack.back()];

			const time_ns duration = time - call.start;
			call.total += duration;

			profile.callstack.pop_back();
		}
	}
}

//
bool save_profile(HSQUIRRELVM vm, const char *path) {
	const auto i = vm_profiles.find(vm);

	if (i == std::end(vm_profiles)) {
		return false;
	}

	const VMProfile &profile = i->second;

	// build a set of all functions (source + name) called during the profile
	// note: funcid is not unique in GS, probably because all scripts use a custom context in which the script is reloaded
	std::set<std::string> ids;

	for (const auto &i : profile.func_info) {
		ids.insert(i.second.source + ":" + i.second.name);
	}

	// compute timings for each function
	std::map<std::string, CallProfile> func_profiles;

	for (const auto &id : ids) {
		std::vector<CallIdx> idxs;

		for (size_t i = 1; i < profile.all_calls.size(); ++i) {
			const Call &call = profile.all_calls[i];

			const FuncInfo &info = profile.func_info.find(call.id.funcid)->second;
			if (id == info.source + ":" + info.name) {
				idxs.push_back(i);
			}
		}

		func_profiles[id] = get_calls_profile(profile.all_calls, idxs);
	}

	// output to CSV
	std::ofstream csv(path);
	csv << "function,total,children,self,hit,ntotal,nself" << std::endl;

	for (const auto &i : func_profiles) {
		csv << i.first << ",";

		const CallProfile &call_profile = i.second;

		const time_ns self = call_profile.total - call_profile.child;

		csv << call_profile.total << ",";
		csv << call_profile.child << ",";
		csv << self << ",";
		csv << call_profile.hit << ",";
		csv << call_profile.total / call_profile.hit << ",";
		csv << self / call_profile.hit;

		csv << std::endl;
	}

	return true;
}

//
SQInteger StartProfiler(HSQUIRRELVM vm) {
	sq_setnativedebughook(vm, native_hook);
	return 0;
}

SQInteger SaveProfile(HSQUIRRELVM vm) {
	__SQ_GETSTART(1)
	__SQ_GETSTRING(path)
	const bool res = save_profile(vm, path);
	__SQ_GETEND
	__SQ_RETURNBOOL(res)
}

SQInteger StopProfiler(HSQUIRRELVM vm) {
	sq_setnativedebughook(vm, nullptr);
	vm_profiles.clear();
	return 0;
}

//
void RegisterProfilerBinding(HSQUIRRELVM vm) {
	GS::Script::sq_register(vm, StartProfiler, "StartProfiler", _SC("."));
	GS::Script::sq_register(vm, SaveProfile, "SaveProfile", _SC(".s"));
	GS::Script::sq_register(vm, StopProfiler, "StopProfiler", _SC("."));
}