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

#include <spslr_list.h>

#include <algorithm>
#include <cstdint>
#include <ostream>
#include <unordered_map>
#include <utility>
#include <vector>

namespace {

struct TARGET_REC {
	uint32_t size;
	uint32_t fieldcnt;
	uint32_t fieldoff;
};

struct TARGET_FIELD_REC {
	uint32_t offset;
	uint32_t size;
	uint32_t alignment;
	uint32_t flags;
};

struct IPIN_REC {
	std::string addr_sym;
	uint32_t size;
	uint32_t program;
};

struct IPIN_OP_REC {
	uint32_t code;
	uint32_t op0;
	uint32_t op1;
};

struct DPIN_REC {
	std::string addr_sym;
	uint32_t target;
};

struct PROGRAM_KEY {
	uint32_t target;
	uint32_t field;

	bool operator==(const PROGRAM_KEY& other) const {
		return target == other.target && field == other.field;
	}
};

struct PROGRAM_KEY_HASH {
	std::size_t operator()(const PROGRAM_KEY& k) const {
		return (static_cast<std::size_t>(k.target) << 32) ^ k.field;
	}
};

static bool emit_header(std::ostream& out);
static bool emit_u32_object(std::ostream& out, const char* name, uint32_t value);
static bool emit_targets(std::ostream& out, const std::vector<TARGET_REC>& targets);
static bool emit_target_fields(std::ostream& out, const std::vector<TARGET_FIELD_REC>& fields);
static bool emit_ipins(std::ostream& out, const std::vector<IPIN_REC>& ipins);
static bool emit_ipin_ops(std::ostream& out, const std::vector<IPIN_OP_REC>& ops);
static bool emit_dpins(std::ostream& out, const std::vector<DPIN_REC>& dpins);

static uint32_t intern_simple_ipin_program(
	std::vector<IPIN_OP_REC>& ops,
	std::unordered_map<PROGRAM_KEY, uint32_t, PROGRAM_KEY_HASH>& memo,
	uint32_t target,
	uint32_t field)
{
	PROGRAM_KEY key{target, field};

	auto it = memo.find(key);
	if (it != memo.end())
		return it->second;

	const uint32_t start = static_cast<uint32_t>(ops.size());

	ops.push_back(IPIN_OP_REC{
		.code = SPSLR_IPIN_OP_ADD_OFFSET,
		.op0  = target,
		.op1  = field,
	});

	ops.push_back(IPIN_OP_REC{
		.code = SPSLR_IPIN_OP_PATCH,
		.op0  = 0,
		.op1  = 0,
	});

	memo.emplace(key, start);
	return start;
}

static bool emit_header(std::ostream& out) {
	out << ".section .spslr,\"a\",@progbits\n";
	out << ".balign 8\n";
	return !!out;
}

static bool emit_u32_object(std::ostream& out, const char* name, uint32_t value) {
	out << ".globl " << name << "\n";
	out << ".type " << name << ", @object\n";
	out << ".balign 4\n";
	out << name << ":\n";
	out << "\t.long " << value << "\n";
	out << ".size " << name << ", 4\n";
	return !!out;
}

static bool emit_targets(std::ostream& out, const std::vector<TARGET_REC>& targets) {
	out << ".globl spslr_targets\n";
	out << ".type spslr_targets, @object\n";
	out << ".balign 4\n";
	out << "spslr_targets:\n";

	for (const TARGET_REC& t : targets) {
		out << "\t.long " << t.size     << "\n";
		out << "\t.long " << t.fieldcnt << "\n";
		out << "\t.long " << t.fieldoff << "\n";
	}

	out << ".size spslr_targets, .-spslr_targets\n";
	return !!out;
}

static bool emit_target_fields(std::ostream& out, const std::vector<TARGET_FIELD_REC>& fields) {
	out << ".globl spslr_target_fields\n";
	out << ".type spslr_target_fields, @object\n";
	out << ".balign 4\n";
	out << "spslr_target_fields:\n";

	for (const TARGET_FIELD_REC& f : fields) {
		out << "\t.long " << f.offset    << "\n";
		out << "\t.long " << f.size      << "\n";
		out << "\t.long " << f.alignment << "\n";
		out << "\t.long " << f.flags     << "\n";
	}

	out << ".size spslr_target_fields, .-spslr_target_fields\n";
	return !!out;
}

static bool emit_ipins(std::ostream& out, const std::vector<IPIN_REC>& ipins) {
	out << ".globl spslr_ipins\n";
	out << ".type spslr_ipins, @object\n";
	out << ".balign 8\n";
	out << "spslr_ipins:\n";

	for (const IPIN_REC& ip : ipins) {
		out << "\t.quad " << ip.addr_sym << "\n";
		out << "\t.long " << ip.size     << "\n";
		out << "\t.long " << ip.program  << "\n";
	}

	out << ".size spslr_ipins, .-spslr_ipins\n";
	return !!out;
}

static bool emit_ipin_ops(std::ostream& out, const std::vector<IPIN_OP_REC>& ops) {
	out << ".globl spslr_ipin_ops\n";
	out << ".type spslr_ipin_ops, @object\n";
	out << ".balign 4\n";
	out << "spslr_ipin_ops:\n";

	for (const IPIN_OP_REC& op : ops) {
		out << "\t.long " << op.code << "\n";
		out << "\t.long " << op.op0  << "\n";
		out << "\t.long " << op.op1  << "\n";
	}

	out << ".size spslr_ipin_ops, .-spslr_ipin_ops\n";
	return !!out;
}

static bool emit_dpins(std::ostream& out, const std::vector<DPIN_REC>& dpins) {
	out << ".globl spslr_dpins\n";
	out << ".type spslr_dpins, @object\n";
	out << ".balign 8\n";
	out << "spslr_dpins:\n";

	for (const DPIN_REC& dp : dpins) {
		out << "\t.quad " << dp.addr_sym << "\n";
		out << "\t.long " << dp.target   << "\n";
	}

	out << ".size spslr_dpins, .-spslr_dpins\n";
	return !!out;
}

}

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

	// Important: target UID == index in spslr_targets[] for the new interface.
	std::vector<TARGET_REC> target_recs(targets.size());
	std::vector<TARGET_FIELD_REC> field_recs;
	field_recs.reserve(64);

	for (uint32_t uid = 0; uid < static_cast<uint32_t>(targets.size()); ++uid) {
		if (!targets.contains(uid))
			return false;

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

		TARGET_REC trec{};
		trec.size = static_cast<uint32_t>(target.size);
		trec.fieldoff = static_cast<uint32_t>(field_recs.size());
		trec.fieldcnt = static_cast<uint32_t>(target.fields.size());

		for (const auto& [off, field] : target.fields) {
			(void)off;
			field_recs.push_back(TARGET_FIELD_REC{
				.offset    = static_cast<uint32_t>(field.offset),
				.size      = static_cast<uint32_t>(field.size),
				.alignment = static_cast<uint32_t>(field.alignment),
				.flags     = static_cast<uint32_t>(field.flags),
			});
		}

		target_recs[uid] = trec;
	}

	std::vector<IPIN_REC> ipin_recs;
	std::vector<IPIN_OP_REC> ipin_ops;
	std::unordered_map<PROGRAM_KEY, uint32_t, PROGRAM_KEY_HASH> program_memo;

	std::vector<DPIN_REC> dpin_recs;

	for (const auto& [cu_uid, cu] : units) {
		(void)cu_uid;

		for (const auto& [sym, ipin] : cu.ipins) {
			const uint32_t global_target = static_cast<uint32_t>(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);

			const uint32_t program = intern_simple_ipin_program(
				ipin_ops,
				program_memo,
				global_target,
				static_cast<uint32_t>(field.idx)
			);

			ipin_recs.push_back(IPIN_REC{
				.addr_sym = ipin.symbol,
				.size     = static_cast<uint32_t>(ipin.imm_size),
				.program  = program,
			});
		}

		for (const auto& [sym, dpin] : cu.dpins) {
			std::vector<DPIN::COMPONENT> sorted_components(dpin.components.begin(), dpin.components.end());

			std::sort(
				sorted_components.begin(),
				sorted_components.end(),
				[](const DPIN::COMPONENT& a, const DPIN::COMPONENT& b) {
					return a.level > b.level;
				}
			);

			for (const DPIN::COMPONENT& component : sorted_components) {
				const uint32_t global_target =
					static_cast<uint32_t>(cu.local_targets.at(component.target));

				std::string addr = dpin.symbol;
				if (component.offset != 0)
					addr += " + " + std::to_string(component.offset);

				dpin_recs.push_back(DPIN_REC{
					.addr_sym = std::move(addr),
					.target   = global_target,
				});
			}
		}
	}

	if (!emit_u32_object(out, "spslr_target_cnt", static_cast<uint32_t>(target_recs.size())))
		return false;
	if (!emit_targets(out, target_recs))
		return false;

	if (!emit_u32_object(out, "spslr_target_field_cnt", static_cast<uint32_t>(field_recs.size())))
		return false;
	if (!emit_target_fields(out, field_recs))
		return false;

	if (!emit_u32_object(out, "spslr_ipin_cnt", static_cast<uint32_t>(ipin_recs.size())))
		return false;
	if (!emit_ipins(out, ipin_recs))
		return false;

	if (!emit_u32_object(out, "spslr_ipin_op_cnt", static_cast<uint32_t>(ipin_ops.size())))
		return false;
	if (!emit_ipin_ops(out, ipin_ops))
		return false;

	if (!emit_u32_object(out, "spslr_dpin_cnt", static_cast<uint32_t>(dpin_recs.size())))
		return false;
	if (!emit_dpins(out, dpin_recs))
		return false;

	return !!out;
}