#include "symbol_collection.h"

#include <iostream>

#define SPSLR_PINPOINT_CU_UID_LABEL "__spslr_cu_" /* suffixed with "<uid>" */

/*
struct GlobalSymtab {
    std::unordered_map<std::string, VADDR> symbols;
};

struct LocalSymtab {
    std::string cu;
    std::unordered_map<std::string, VADDR> symbols;
};
*/

using namespace LIEF::ELF;

GlobalSymtab global_syms;
std::unordered_map<std::string, LocalSymtab> local_syms;

bool collect_symbols(const std::unique_ptr<LIEF::ELF::Binary>& bin) {
	if (!bin)
		return false;

	bool cu_uid_found = false;
	LocalSymtab tmp_local_syms;

	auto finish_cu_symbols = [&]() -> bool {
		if (cu_uid_found) {
			if (local_syms.contains(tmp_local_syms.cu)) {
				std::cerr << "Duplicate CU UID -> " << tmp_local_syms.cu << std::endl;
				return false;
			}

			local_syms.emplace(tmp_local_syms.cu, std::move(tmp_local_syms));
		}

		tmp_local_syms.cu.clear();
		tmp_local_syms.symbols.clear();
		cu_uid_found = false;
		return true;
	};

	for (const Symbol& sym : bin->symbols()) {
		if (sym.type() == Symbol::TYPE::FILE) {
			if (!finish_cu_symbols())
				return false;
		}

		if (sym.name().empty())
			continue;

		if (sym.type() != Symbol::TYPE::OBJECT && sym.type() != Symbol::TYPE::NOTYPE)
			continue;

		VADDR vaddr = sym.value();

		if (sym.binding() == Symbol::BINDING::GLOBAL) {
			if (global_syms.symbols.contains(sym.name())) {
				std::cerr << "Duplicate global symbol \"" << sym.name() << "\"!" << std::endl;
				return false;
			}

			global_syms.symbols.emplace(sym.name(), vaddr);
			continue;
		}

		if (sym.binding() != Symbol::BINDING::LOCAL)
			continue;

		if (sym.name().starts_with(SPSLR_PINPOINT_CU_UID_LABEL)) {
			tmp_local_syms.cu = sym.name().substr(std::strlen(SPSLR_PINPOINT_CU_UID_LABEL));
			cu_uid_found = true;
			continue;
		}

		if (tmp_local_syms.symbols.contains(sym.name())) {
			std::cerr << "Duplicate local symbol \"" << sym.name() << "\"!" << std::endl;
			return false;
		}

		tmp_local_syms.symbols.emplace(sym.name(), vaddr);
		continue;
	}

	if (!finish_cu_symbols())
		return false;

	return true;
}