#include <psemek/util/profiler.hpp>
#include <psemek/util/pretty_print.hpp>
#include <psemek/util/statistics.hpp>
#include <psemek/log/log.hpp>

#include <map>
#include <unordered_map>
#include <memory>
#include <iomanip>
#include <thread>
#include <mutex>
#include <vector>

namespace psemek::util
{

	namespace
	{

		struct profiler_tree
		{
			profiler_tree * parent = nullptr;
			std::mutex * mutex = nullptr;

			statistics<double> execution_time;

			std::map<std::string, profiler_tree> children;
		};

		static std::map<std::thread::id, profiler_tree> roots;
		static std::mutex roots_mutex;
		static std::unordered_map<std::thread::id, std::mutex> thread_mutex;

		static thread_local profiler_tree * current = nullptr;

		void dump_impl(profiler_tree const & node, std::size_t depth)
		{
			std::size_t max_name_length = 0;

			for (auto const & c : node.children)
				max_name_length = std::max(max_name_length, c.first.size());

			for (auto const & c : node.children)
			{
				log::info() << std::string(depth * 2, ' ') << std::setw(max_name_length + 3) << std::left << std::setfill(' ') << c.first
					<< util::pretty(std::chrono::duration<double>{c.second.execution_time.mean()}, std::chrono::microseconds{})
					<< " (" << c.second.execution_time.count() << " calls)";
				dump_impl(c.second, depth + 1);
			}
		}

		profiler_tree merge(profiler_tree t1, profiler_tree const & t2)
		{
			t1.execution_time = merge(t1.execution_time, t2.execution_time);

			for (auto const & p : t2.children)
			{
				auto it = t1.children.find(p.first);
				if (it == t1.children.end())
				{
					t1.children.insert(p);
				}
				else
				{
					it->second = merge(std::move(it->second), p.second);
				}
			}

			return t1;
		}

		profiler_tree merged_tree()
		{
			std::vector<profiler_tree const *> r;

			{
				std::lock_guard lock{roots_mutex};
				for (auto const & p : roots)
					r.push_back(&p.second);
			}

			profiler_tree result;

			for (auto root : r)
			{
				std::lock_guard lock{*(root->mutex)};
				result = merge(std::move(result), *root);
			}

			return result;
		}

	}

	profiler::profiler(std::string const & name)
	{
		auto const id = std::this_thread::get_id();

		if (!current)
		{
			std::lock_guard lock{roots_mutex};
			current = &roots[id];
			current->mutex = &thread_mutex[id];
		}

		std::lock_guard lock{*(current->mutex)};

		auto next = &current->children[name];
		next->parent = current;
		next->mutex = current->mutex;
		current = next;
		clock_.restart();
	}

	profiler::~profiler()
	{
		std::lock_guard lock{*(current->mutex)};
		current->execution_time.push(clock_.count());
		current = current->parent;
	}

	void profiler::dump()
	{
		dump_impl(merged_tree(), 0);
	}

}