#include <pslang/jit/linker.hpp>
#include <pslang/jit/resolver.hpp>

namespace pslang::jit
{

	namespace
	{

		std::size_t align(std::size_t value, std::size_t alignment)
		{
			return ((value + alignment - 1) / alignment) * alignment;
		}

	}

	linked_program_info prepare_link(std::vector<compiled_module> const & modules, link_options const & options)
	{
		linked_program_info result;

		std::size_t total_code_size = 0;
		std::size_t total_data_size = 0;
		std::size_t total_relo_size = options.reserved_relocations * 8;

		for (auto const & module : modules)
		{
			for (auto const & function : module.functions)
				result.functions[function.first] = total_code_size + function.second;

			for (auto const & global : module.globals)
				result.globals[global.first] = total_data_size + global.second;

			for (auto const & relocation : module.external_function_resolve)
				if (!result.relocation_offset.contains(relocation.name))
				{
					result.relocation_offset[relocation.name] = total_relo_size;
					total_relo_size += 8;
				}

			if (module.entry_point)
				result.entry_points.push_back(total_code_size + module.functions.at(*module.entry_point));
			else
				result.entry_points.push_back(std::nullopt);

			total_code_size += module.code.size();
			total_data_size += module.data.size();
		}

		total_code_size = align(total_code_size, options.section_alignment);
		total_data_size = align(total_data_size, options.section_alignment);
		total_relo_size = align(total_relo_size, options.section_alignment);

		std::size_t const data_offset = total_code_size;
		std::size_t const relo_offset = total_code_size + total_data_size;
		result.total_size = total_code_size + total_data_size + total_relo_size;

		for (auto & global : result.globals)
			global.second += data_offset;

		for (auto & relocation : result.relocation_offset)
			relocation.second += relo_offset;

		result.code = {0, total_code_size};
		result.data = {total_code_size, total_data_size};
		result.relocations = {total_code_size + total_data_size, total_relo_size};

		return result;
	}

	void link(std::vector<compiled_module> const & modules, linked_program_info const & info, std::uint8_t * storage)
	{
		// TODO: resolve internal symbols (across different modules)

		if (modules.empty())
			return;

		auto resolver = make_resolver(modules.front().abi.isa);

		{
			auto dst = storage + info.code.begin;

			for (auto const & module : modules)
			{
				auto end = std::copy(module.code.data(), module.code.data() + module.code.size(), dst);

				for (auto const & resolve : module.internal_global_resolve)
				{
					std::int32_t offset = info.globals.at(resolve.global) - (dst - storage + resolve.instruction_offset);
					resolver->resolve(dst + resolve.instruction_offset, offset);
				}

				for (auto const & resolve : module.external_function_resolve)
				{
					std::int32_t offset = info.relocation_offset.at(resolve.name) - (dst - storage + resolve.instruction_offset);
					resolver->resolve(dst + resolve.instruction_offset, offset);
				}

				dst = end;
			}
		}

		{
			auto dst = storage + info.data.begin;

			for (auto const & module : modules)
				dst = std::copy(module.data.data(), module.data.data() + module.data.size(), dst);
		}
	}

}