#include <iostream>
#include <final.h>
#include <stage0.h>
#include <stage1.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 bool src_dir_known = false, dst_dir_known = false;
static std::string src_dir, dst_dir;

void set_src_dir(const char* dir) {
	if (!dir) {
		src_dir_known = false;
		return;
	}

	src_dir = std::string{ dir };
	src_dir_known = true;
}

void set_dst_dir(const char* dir) {
	if (!dir) {
		dst_dir_known = false;
		return;
	}

	dst_dir = std::string{ dir };
	dst_dir_known = true;
}

bool has_src_dir() {
	return src_dir_known;
}

bool has_dst_dir() {
	return dst_dir_known;
}

const char* get_src_dir() {
	if (!src_dir_known)
		return nullptr;

	return src_dir.c_str();
}

const char* get_dst_dir() {
	if (!dst_dir_known)
		return nullptr;

	return dst_dir.c_str();
}

static std::filesystem::path relative_src_path() {
	const char* src_root = get_src_dir();
	if (!src_root)
		pinpoint_fatal("relative_src_path is missing the source root path");

	if (!main_input_filename)
		pinpoint_fatal("relative_src_path is missing the current source file path");

	std::filesystem::path abs_current_src = std::filesystem::absolute(main_input_filename);
	std::filesystem::path abs_root_src = std::filesystem::absolute(src_root);
	return std::filesystem::relative(abs_current_src, abs_root_src);
}

static std::filesystem::path spslr_output_file(const std::filesystem::path& infile) {
	const char* dst_root = get_dst_dir();
	if (!dst_root)
		pinpoint_fatal("spslr_output_file is missing the destination root path");

	std::filesystem::path abs_root_dst = std::filesystem::absolute(dst_root);
	return (abs_root_dst / infile).string() + SPSLR_PINFILE_EXTENSION;
}

static std::string calculate_cu_uid(const std::filesystem::path& infile) {
	unsigned char md5_digest[16];

	const char* infile_cstr = infile.c_str();
	md5_buffer(infile_cstr, strlen(infile_cstr), md5_digest);

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

	return std::string{ md5_digest_hex };
}

static void emit_cu_uid_label(const std::string& uid) {
	char label[128];
	snprintf(label, sizeof(label), SPSLR_PINPOINT_CU_UID_LABEL "%s", uid.c_str());

	fprintf(asm_out_file, "%s:\n", label);
}

static std::ofstream open_spslr_output_file(const std::filesystem::path& p) {
	std::filesystem::create_directories(p.parent_path());
	std::ofstream out(p);
	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::filesystem::path infile = relative_src_path();
	std::filesystem::path outfile = spslr_output_file(infile);

	std::string cu_uid = calculate_cu_uid(infile);

	emit_cu_uid_label(cu_uid);

	// Dump all accumulated data to spslr file

	std::ofstream out = open_spslr_output_file(outfile);

	// Header associates data with compilation unit

	out << "SPSLR " << infile.string() << " " << 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>
			std::size_t field_alignment = field.size; // TODO
			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 <local/global> <symbol> <offset> <level> <target uid>
			out << "dpin " << (dpin.global ? "g" : "l") << " " << dpin.symbol << " "
				<< c.offset << " " << c.level << " " << c.target << std::endl;
		}
	}

	// Dump all instruction pins

	for (const auto& [uid, ipin] : S1InstructionPin::all()) {
		// ipin <label> <target uid> <field offset>
		out << "ipin " << SPSLR_PINPOINT_STAGE1_PIN << uid << " "
			<< ipin.target << " " << ipin.offset << std::endl;
	}
}