#include <iostream>
#include <final.h>
#include <stage0.h>
#include <stage2.h>
#include <string>
#include <filesystem>
#include <fstream>

#include <safe-input.h>
#include <safe-output.h>
#include <safe-md5.h>
#include <pinpoint_config.h>
#include <pinpoint_error.h>

static std::string src_file, cu_hash;

static bool meta_dir_known = false, srcroot_dir_known = false;
static std::string meta_dir, srcroot_dir;

bool init_src_file() {
	if (!main_input_filename || !has_srcroot_dir())
		return false;

	namespace fs = std::filesystem;

	// Canonicalize both paths (robust against symlinks, ../, etc.)
	fs::path src_abs = fs::weakly_canonical(fs::absolute(main_input_filename));
	fs::path root_abs = fs::weakly_canonical(fs::absolute(get_srcroot_dir()));

	// Ensure src_abs is actually inside root_abs
	auto src_it = src_abs.begin();
	auto root_it = root_abs.begin();

	for (; root_it != root_abs.end(); ++root_it, ++src_it) {
		if (src_it == src_abs.end() || *src_it != *root_it) {
			// Not under srcroot
			std::cerr << "The source file " << src_abs << " is not under source root " << root_abs << std::endl;
			return false;
		}
	}

	// Compute relative path
	fs::path rel = fs::relative(src_abs, root_abs);

	// Normalize to forward slashes (important for metadata consistency)
	src_file = rel.generic_string();

	return true;
}

const char* get_src_file() {
	return src_file.c_str();
}

bool init_cu_hash() {
	unsigned char md5_digest[16];
	const char* s = get_src_file();
	md5_buffer(s, strlen(s), md5_digest);

	char md5_digest_hex[33] = {};
	for (int i = 0; i < 16; i++)
		sprintf(md5_digest_hex + i * 2, "%02x", md5_digest[i]);

	cu_hash = std::string{md5_digest_hex};
	return true;
}

const char* get_cu_hash() {
	return cu_hash.c_str();
}

void set_meta_dir(const char* dir) {
	if (!dir) {
		meta_dir_known = false;
		return;
	}

	meta_dir = std::string{ dir };
	meta_dir_known = true;
}

bool has_meta_dir() {
	return meta_dir_known;
}

const char* get_meta_dir() {
	if (!meta_dir_known)
		return nullptr;

	return meta_dir.c_str();
}

void set_srcroot_dir(const char* dir) {
	if (!dir) {
		srcroot_dir_known = false;
		return;
	}

	srcroot_dir = std::string{ dir };
	srcroot_dir_known = true;
}

bool has_srcroot_dir() {
	return srcroot_dir_known;
}

const char* get_srcroot_dir() {
	if (!srcroot_dir_known)
		return nullptr;

	return srcroot_dir.c_str();
}

static void emit_dpin_alias_labels() {
	for (const DataPin& dpin : DataPin::all()) {
		if (dpin.pin_symbol.empty() || dpin.symbol.empty())
			pinpoint_fatal("emit_dpin_alias_labels got incomplete data pin");

		fprintf(asm_out_file, ".globl %s\n", dpin.pin_symbol.c_str());
		fprintf(asm_out_file, ".set %s, %s\n", dpin.pin_symbol.c_str(), dpin.symbol.c_str());
	}
}

static std::ofstream open_spslr_output_file() {
	std::filesystem::path metadir { get_meta_dir() };
	std::filesystem::path metafile { std::string{ get_src_file() } + SPSLR_PINFILE_EXTENSION };

	std::filesystem::path op = metadir / metafile;

	std::filesystem::create_directories(op.parent_path());

	std::cout << "Dumping meta data to " << op << std::endl;

	std::ofstream out(op);
	if (!out)
		pinpoint_fatal("open_spslr_output_file failed to open spslr dump file");

	return std::move(out);
}

void on_finish_unit(void* plugin_data, void* user_data) {
	std::string cu_uid = get_cu_hash();

	// Emit globally unique data pin label for each static object to be randomized

	emit_dpin_alias_labels();

	// Dump all accumulated data to spslr file

	std::ofstream out = open_spslr_output_file();

	// Header associates data with compilation unit

	out << "SPSLR " << get_src_file() << " " << cu_uid << std::endl;

	// Dump all target structs

	for (const auto& [uid, target] : TargetType::all()) {
		// target <name> <local uid> <size> <field count>
		out << "target " << target.name() << " " << uid << " " << target.size()
			<< " " << target.fields().size() << std::endl;

		if (!target.has_fields())
			pinpoint_fatal("on_finish_unit encountered incomplete target type");

		for (const auto& [off, field] : target.fields()) {
			// f <offset> <size> <alignment> <flags>
			out << "f " << field.offset << " " << field.size << " " << field.alignment << " " << field.flags << std::endl;
		}
	}

	// Dump all data pins

	for (const DataPin& dpin : DataPin::all()) {
		for (const DataPin::Component& c : dpin.components) {
			// dpin <symbol> <offset> <level> <target uid>
			out << "dpin " << " " << dpin.pin_symbol << " "
				<< c.offset << " " << c.level << " " << c.target << std::endl;
		}
	}

	// Dump all instruction pins

	for (const auto& [uid, ipin] : s2_pins()) {
		// ipin <symbol> <target uid> <field offset> <immediate size>
		out << "ipin " << ipin.symbol << " "
			<< ipin.target << " " << ipin.offset << " " << ipin.imm_size << std::endl;
	}
}