#include "module.h"

#include "dbg.h"
#include "filesystem.h"
#include "strtools.h"

#include <elf.h>
#include <fcntl.h>
#include <link.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

namespace modules
{
	CModule* engine = nullptr;
	CModule* server = nullptr;
	CModule* schemasystem = nullptr;
}

static bool GetModuleInformation(void* handle, void** base, size_t* length, std::vector<Section>& sections)
{
	link_map* lmap = nullptr;
	if (dlinfo(handle, RTLD_DI_LINKMAP, &lmap) != 0 || !lmap)
		return false;

	int fd = open(lmap->l_name, O_RDONLY);
	if (fd == -1)
		return false;

	struct stat st {};
	if (fstat(fd, &st) != 0)
	{
		close(fd);
		return false;
	}

	void* map = mmap(nullptr, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
	close(fd);
	if (map == MAP_FAILED)
		return false;

	auto* ehdr = static_cast<ElfW(Ehdr)*>(map);
	auto* shdrs = reinterpret_cast<ElfW(Shdr)*>(reinterpret_cast<uintptr_t>(ehdr) + ehdr->e_shoff);
	const char* strTab = reinterpret_cast<const char*>(reinterpret_cast<uintptr_t>(ehdr) + shdrs[ehdr->e_shstrndx].sh_offset);

	for (int i = 0; i < ehdr->e_phnum; ++i)
	{
		auto* phdr = reinterpret_cast<ElfW(Phdr)*>(reinterpret_cast<uintptr_t>(ehdr) + ehdr->e_phoff + i * ehdr->e_phentsize);
		if (phdr->p_type == PT_LOAD && (phdr->p_flags & PF_X))
		{
			*base = reinterpret_cast<void*>(lmap->l_addr + phdr->p_vaddr);
			*length = phdr->p_filesz;
			break;
		}
	}

	for (int i = 0; i < ehdr->e_shnum; ++i)
	{
		auto* shdr = reinterpret_cast<ElfW(Shdr)*>(reinterpret_cast<uintptr_t>(shdrs) + i * ehdr->e_shentsize);
		if (*(strTab + shdr->sh_name) == '\0')
			continue;
		sections.push_back({strTab + shdr->sh_name, reinterpret_cast<void*>(lmap->l_addr + shdr->sh_addr), shdr->sh_size});
	}

	munmap(map, st.st_size);
	return *base && *length;
}

CModule::CModule(const char* relativePath, const char* module) : m_name(module)
{
	char path[MAX_PATH];
	V_snprintf(path, sizeof(path), "%s%s%s%s%s", Plat_GetGameDirectory(), relativePath, MODULE_PREFIX, module, MODULE_EXT);

	m_handle = dlopen(path, RTLD_NOW | RTLD_NOLOAD);
	if (!m_handle)
		m_handle = dlopen(path, RTLD_NOW);
	if (!m_handle)
		Error("[SendProxy] Could not open %s: %s\n", path, dlerror());

	if (!GetModuleInformation(m_handle, &m_base, &m_size, m_sections))
		Error("[SendProxy] Could not inspect %s\n", path);
}

void* CModule::FindSignature(const uint8_t* pattern, size_t length, int& error) const
{
	error = SIG_NOT_FOUND;
	void* found = nullptr;
	auto* memory = static_cast<uint8_t*>(m_base);

	for (size_t i = 0; i + length < m_size; ++i)
	{
		size_t matched = 0;
		while (matched < length && (pattern[matched] == 0x2A || memory[i + matched] == pattern[matched]))
			++matched;
		if (matched != length)
			continue;
		if (found)
		{
			error = SIG_FOUND_MULTIPLE;
			return found;
		}
		found = memory + i;
		error = SIG_OK;
	}

	return found;
}

void* CModule::FindInterface(const char* name) const
{
	auto fn = reinterpret_cast<CreateInterfaceFn>(dlsym(m_handle, "CreateInterface"));
	if (!fn)
		return nullptr;
	return fn(name, nullptr);
}