#include <psemek/prof/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::prof
{

	namespace
	{

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

			util::statistics<double> execution_time;

			std::map<std::string, profiler_tree> children;
			std::vector<std::map<std::string, profiler_tree>::iterator> children_list;

			profiler_tree * child(std::string const & name)
			{
				auto it = children.find(name);
				if (it == children.end())
				{
					it = children.insert({name, profiler_tree{}}).first;
					children_list.push_back(it);
				}
				return &(it->second);
			}
		};

		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;

		struct pretty_impl
		{
			double value;
		};

		pretty_impl pretty(double value)
		{
			return {value};
		}

		std::ostream & operator << (std::ostream & os, pretty_impl p)
		{
			char const * suf[] =
			{
				"ns",
				"us",
				"ms",
				"s",
			};

			double s = 1e-9;

			for (std::size_t i = 0; i < 4; ++i)
			{
				if (p.value < (s * 1000.0) || i == 3)
				{
					os << (p.value / s) << ' ' << suf[i];
					return os;
				}

				s *= 1000.0;
			}

			return os;
		}

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

			for (auto it : node.children_list)
				max_name_length = std::max(max_name_length, it->first.size());

			for (auto it : node.children_list)
			{
				auto const & c = *it;
				auto const & stat = c.second.execution_time;
				if (stat.count() == 1)
					log::info() << std::string(depth * 2, ' ') << std::setw(max_name_length + 3) << std::left << std::setfill(' ') << c.first
								<< " " << pretty(stat.mean())
								<< " (" << stat.count() << " call)";
				else
					log::info() << std::string(depth * 2, ' ') << std::setw(max_name_length + 3) << std::left << std::setfill(' ') << c.first
						<< " avg " << pretty(stat.mean())
						<< "    25% " << pretty(stat.percentile(0.25))
						<< "    50% " << pretty(stat.percentile(0.50))
						<< "    75% " << pretty(stat.percentile(0.75))
						<< " (" << stat.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 it2 : t2.children_list)
			{
				auto it1 = t1.children.find(it2->first);
				if (it1 == t1.children.end())
				{
					it1 = t1.children.insert(*it2).first;
					t1.children_list.push_back(it1);
				}
				else
				{
					it1->second = merge(std::move(it1->second), it2->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->child(name);
		next->parent = current;
		next->mutex = current->mutex;
		current = next;
		clock_.restart();
	}

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

	void profiler::push(std::string const & name, std::chrono::duration<double> duration)
	{
		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 child = current->child(name);
		child->parent = current;
		child->mutex = current->mutex;
		child->execution_time.push(duration.count());
	}

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

}