#include "emit.h"
#include "accumulation.h"

#include <spslr_program.h>

static bool emit_header(std::ostream& out);
static bool emit_target_header(std::ostream& out, std::size_t uid, std::size_t size, std::size_t fieldcnt);
static bool emit_target_field(std::ostream& out, std::size_t offset, std::size_t size, std::size_t alignment, std::size_t flags);
static bool emit_target_randomize(std::ostream& out, std::size_t uid);
static bool emit_target(std::ostream& out, std::size_t uid, const TARGET& target);
static bool emit_ipin(std::ostream& out, const CU& cu, const IPIN& ipin);
static bool emit_dpin_component(std::ostream& out, std::size_t target, const std::string& sym, std::size_t offset);
static bool emit_dpin(std::ostream& out, const CU& cu, const DPIN& dpin);
static bool emit_exit(std::ostream& out);

bool emit_patcher_program_asm(std::ostream& out) {
	if (!emit_header(out))
		return false;

	// Dump all targets
	for (const auto& [uid, target] : targets) {
		if (!emit_target(out, uid, target))
			return false;
	}

	// Dump all pins
	for (const auto& [uid, cu] : units) {
		for (const auto& [sym, pin] : cu.ipins) {
			if (!emit_ipin(out, cu, pin))
				return false;
		}

		for (const auto& [sym, pin] : cu.dpins) {
			if (!emit_dpin(out, cu, pin))
				return false;
		}
	}

	// Exit patcher program
	if (!emit_exit(out))
		return false;

	return true;
}

bool emit_header(std::ostream& out) {
	out << ".section .spslr,\"a\",@progbits\n";
	out << ".balign 1\n";
	out << ".globl __spslr_program\n";
	out << ".type __spslr_program, @object\n";
	out << "__spslr_program:\n";
	return !!out;
}

bool emit_target_header(std::ostream& out, std::size_t uid, std::size_t size, std::size_t fieldcnt) {
	out << "\t.byte " << OPCODE_SPSLR_TARGET << "\n";
	out << "\t.long " << uid << "\n";
	out << "\t.long " << size << "\n";
	out << "\t.long " << fieldcnt << "\n";
	return !!out;
}

bool emit_target_field(std::ostream& out, std::size_t offset, std::size_t size, std::size_t alignment, std::size_t flags) {
	out << "\t.byte " << OPCODE_SPSLR_FIELD << "\n";
	out << "\t.long " << offset << "\n";
	out << "\t.long " << size << "\n";
	out << "\t.long " << alignment << "\n";
	out << "\t.long " << flags << "\n";
	return !!out;
}

bool emit_target_randomize(std::ostream& out, std::size_t uid) {
	out << "\t.byte " << OPCODE_SPSLR_RANDOMIZE << "\n";
	out << "\t.long " << uid << "\n";
	return !!out;
}

bool emit_target(std::ostream& out, std::size_t uid, const TARGET& target) {
	if (!emit_target_header(out, uid, target.size, target.fields.size()))
		return false;

	for (const auto& [off, field] : target.fields) {
		// Note -> Might want to do explicit flag mapping from pinpoint to selfpatch
		if (!emit_target_field(out, field.offset, field.size, field.alignment, field.flags))
			return false;
	}

	if (!emit_target_randomize(out, uid))
		return false;

	return true;
}

bool emit_ipin(std::ostream& out, const CU& cu, const IPIN& ipin) {
	std::size_t global_target = cu.local_targets.at(ipin.local_target);

	if (!targets.contains(global_target))
		return false;

	const TARGET& target = targets.at(global_target);

	if (!target.fields.contains(ipin.field_offset))
		return false;

	const FIELD& field = target.fields.at(ipin.field_offset);

	out << "\t.byte " << OPCODE_SPSLR_IPATCH << "\n";
	out << "\t.quad " << ipin.symbol << "\n";
	out << "\t.long " << ipin.imm_size << "\n";
	out << "\t.long " << global_target << "\n";
	out << "\t.long " << field.idx << "\n";
	return !!out;
}

bool emit_dpin_component(std::ostream& out, std::size_t target, const std::string& sym, std::size_t offset) {
	out << "\t.byte " << OPCODE_SPSLR_DPATCH << "\n";
	out << "\t.quad " << sym;
	if (offset != 0)
		out << " + " << offset;
	out << "\n";
	out << "\t.long " << target << "\n";
	return !!out;
}

bool emit_dpin(std::ostream& out, const CU& cu, const DPIN& dpin) {
	std::list<DPIN::COMPONENT> sorted_components = dpin.components;
	sorted_components.sort([](const DPIN::COMPONENT& a, const DPIN::COMPONENT& b) {
		return a.level > b.level; // Descending order to handle deeper nested first
	});

	for (const DPIN::COMPONENT& component : sorted_components) {
		std::size_t global_target = cu.local_targets.at(component.target);
		if (!emit_dpin_component(out, global_target, dpin.symbol, component.offset))
			return false;
	}

	return true;
}

bool emit_exit(std::ostream& out) {
	out << "\t.byte " << OPCODE_SPSLR_EXIT << "\n";
	return !!out;
}