#include <iostream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <map>
#include <list>

#include <gcc-plugin.h>
#include <plugin-version.h>
#include <tree.h>
#include <langhooks.h>
#include <gimple.h>
#include <gimplify.h>
#include <stringpool.h>
#include <tree-pass.h>
#include <gimple-iterator.h>
#include <gimple-pretty-print.h>
#include <tree-pretty-print.h>
#include <tree-dump.h>
#include <print-tree.h>
#include <tree-iterator.h>
#include <rtl.h>
#include <memmodel.h>
#include <emit-rtl.h>
#include <df.h>

int plugin_is_GPL_compatible;

#ifndef SPSLR_OFFSETOF
#define SPSLR_OFFSETOF "__spslr_offsetof_" /* with suffix <uid> */
#endif

#ifndef UNSPEC_SPSLR_OFFSETOF
#define UNSPEC_SPSLR_OFFSETOF 1042
#endif

// At early RTL, replace __spslr_offsetof_<uid> with UNSPECs to avoid ABI and clobbering "problems"

static bool extract_callee_symbol (rtx call_rtx, const char** out_name) {
	if (!call_rtx || GET_CODE(call_rtx) != CALL)
		return false;

	rtx op0 = XEXP(call_rtx, 0);
	if (!op0)
		return false;

	rtx addr = XEXP(op0, 0);
	if (addr && GET_CODE(addr) == SYMBOL_REF) {
		*out_name = XSTR(addr, 0);
		return true;
	}

	return false;
}

static bool parse_uid_from_name(const char* name, unsigned long* out_uid) {
	if (!name || !out_uid)
		return false;

	if (strncmp(name, SPSLR_OFFSETOF, strlen(SPSLR_OFFSETOF)) != 0)
		return false;

	const char *p = name + strlen(SPSLR_OFFSETOF);
	if (*p == 0)
		return false;

	char *endp = nullptr;
	unsigned long v = strtoul(p, &endp, 10);
	if (endp == p)
		return false;

	*out_uid = v;
	return true;
}

static bool extract_dest_and_call (rtx pat, rtx* out_dest, rtx* out_call) {
	if (!pat)
		return false;

	if (GET_CODE(pat) == SET) {
		rtx src = SET_SRC(pat);
		if (GET_CODE(src) == CALL) {
			*out_dest = SET_DEST(pat);
			*out_call = src;
			return true;
		}
		return false;
	}

	if (GET_CODE(pat) == PARALLEL) {
		// Look for a SET whose src is CALL
		int n = XVECLEN(pat, 0);
		for (int i = 0; i < n; ++i) {
			rtx elt = XVECEXP(pat, 0, i);
			if (GET_CODE(elt) == SET && GET_CODE(SET_SRC(elt)) == CALL) {
				*out_dest = SET_DEST(elt);
				*out_call = SET_SRC(elt);
				return true;
			}
		}
	}

	return false;
}

static void call_to_unspec(function* fn) {
	if (!fn)
		return;

	unsigned replaced = 0;

	basic_block bb;
	FOR_EACH_BB_FN(bb, fn) {
		for (rtx_insn* insn = BB_HEAD(bb); insn != NEXT_INSN(BB_END(bb)); insn = NEXT_INSN(insn)) {
			if (!INSN_P(insn) || !CALL_P(insn))
				continue;

			rtx pat = PATTERN(insn);
			rtx dest = NULL_RTX;
			rtx call = NULL_RTX;
			if (!extract_dest_and_call(pat, &dest, &call))
				continue;

			const char *name = nullptr;
			if (!extract_callee_symbol(call, &name))
				continue;

			unsigned long uid = 0;
			if (!parse_uid_from_name(name, &uid))
				continue;

			// We expect a returning call (assigned to dest)
			if (!dest || !REG_P(dest) || GET_MODE(dest) == VOIDmode)
				continue;

			machine_mode mode = GET_MODE(dest);

			// Build: (set dest (unspec:mode [(const_int uid)] UNSPEC_SPSLR_OFFSETOF))
			rtvec vec = gen_rtvec(1, GEN_INT((HOST_WIDE_INT) uid));
			rtx uns = gen_rtx_UNSPEC(mode, vec, UNSPEC_SPSLR_OFFSETOF); // Note -> maybe volatile
			rtx set = gen_rtx_SET(dest, uns);

			emit_insn_before(set, insn);
			delete_insn(insn);
			replaced++;
		}
	}

	if (replaced) {
		df_set_bb_dirty(ENTRY_BLOCK_PTR_FOR_FN(fn));
		df_set_bb_dirty(EXIT_BLOCK_PTR_FOR_FN(fn));
	}
}

const pass_data call_to_unspec_pass_data = {
	RTL_PASS,
	"call_to_unspec",
	OPTGROUP_NONE,
	TV_NONE,
	PROP_rtl,
	0, 0, 0, 0
};

struct call_to_unspec_pass : rtl_opt_pass {
	call_to_unspec_pass(gcc::context* ctxt)
		: rtl_opt_pass(call_to_unspec_pass_data, ctxt) {}

	unsigned int execute(function* fn) override {
		call_to_unspec(fn);
		return 0;
	}
};

// Late RTL pass replaces UNSPECs with labeled constants (must happen before vregs, no optimizations afterwards)

static bool lookup_initial_member_offset(SPSLROffsetofCallData::UID uid, Member::OFF& ioff) {
	auto call_data_it = spslr_offsetof_calls.find(uid);
	if (call_data_it == spslr_offsetof_calls.end())
		return false;

	const SPSLROffsetofCallData& call_data = call_data_it->second;

	// call_data.target, call_data.member

	ioff = call_data.member;
	return true;
}

static void emit_named_asm_label_before(SPSLROffsetofCallData::UID uid, rtx_insn* before) {
	char name[128];
	snprintf(name, sizeof(name), "%s%lu:\n", SPSLR_OFFSETOF, uid);

	rtvec no_out = rtvec_alloc(0);
	rtvec no_in = rtvec_alloc(0);
	rtvec no_cl = rtvec_alloc(0);
	const char* empty_constraints = "";

	location_t loc = INSN_LOCATION(before);

	rtx asmops = gen_rtx_ASM_OPERANDS (VOIDmode, ggc_strdup(name), empty_constraints, 1, no_out, no_in, no_cl, loc);
	emit_insn_before (asmops, before);
}

static rtx labeled_cst_mov(SPSLROffsetofCallData::UID uid, rtx dest, Member::OFF ioff, location_t loc) {
	char asm_str[128];
	snprintf(asm_str, sizeof(asm_str), "%s%lu:\nmov %1, %0\n", SPSLR_OFFSETOF, uid);

	// rtl.def
	// -> DEF_RTL_EXPR(ASM_INPUT, "asm_input", "sL", RTX_EXTRA) -> only string+location in ASM_INPUT
	// -> DEF_RTL_EXPR(ASM_OPERANDS, "asm_operands", "ssiEEEL", RTX_EXTRA)

	rtx desc_in1 = gen_rtx_ASM_INPUT(SImode, ggc_strdup("i"));
	rtvec desc_inputs = gen_rtvec(1, desc_in1);
	rtvec inputs = gen_rtvec(1, GEN_INT(ioff));

	const char* desc_outputs = "=r";
	rtvec outputs = gen_rtvec(1, dest);

	rtvec labels = rtvec_alloc(0);

	rtx asmops = gen_rtx_ASM_OPERANDS(GET_MODE(dest),
		ggc_strdup(asm_str),      /* template */
		ggc_strdup(desc_outputs), /* output constraint */
		0,                        /* output number */
		inputs,                   /* vector of input RTXs */
		desc_inputs,              /* vector of input descriptors */
		labels,                   /* labels (empty) */
		loc);                     /* source location */

	rtx cc_clob = gen_rtx_CLOBBER(VOIDmode, gen_rtx_REG(CCmode, 17));

	rtvec vec = gen_rtvec (2, gen_rtx_SET(dest, asmops), cc_clob);
	rtx parallel = gen_rtx_PARALLEL(VOIDmode, vec);

	return parallel;
}

static void unspec_to_labeled_const(function* fn) {
	basic_block bb;
	FOR_EACH_BB_FN(bb, fn) {
		for (rtx_insn *insn = BB_HEAD(bb); insn != NEXT_INSN(BB_END(bb)); insn = NEXT_INSN(insn)) {
			if (!INSN_P(insn))
				continue;

			rtx pat = PATTERN(insn);
			if (GET_CODE(pat) != SET)
				continue;

			rtx src = SET_SRC(pat);
			if (GET_CODE(src) != UNSPEC)
				continue;

			if (XINT(src, 1) != UNSPEC_SPSLR_OFFSETOF)
				continue;

			/* Extract UID from UNSPEC operands.  */
			if (XVEC(src, 0) == NULL || XVECLEN(src, 0) < 1)
				continue;

			rtx arg = XVECEXP(src, 0, 0);
			if (!CONST_INT_P(arg))
				continue;

			SPSLROffsetofCallData::UID uid = (SPSLROffsetofCallData::UID)INTVAL(arg);
			rtx dest = SET_DEST(pat);

			Member::OFF initial_offset = uid; // TODO -> Use value that forces 32 bit (currently uses 32 bit anyways)?
			if (!lookup_initial_member_offset(uid, initial_offset)) {
				std::cerr << "Failed to query initial member offset for access uid " << uid << "!" << std::endl;
				return;
			}

			// TODO

			PATTERN(insn) = labeled_cst_mov(uid, dest, initial_offset, INSN_LOCATION(insn));
			INSN_CODE(insn) = -1;
			df_insn_rescan(insn);

			/*
			// Generate asm label
			emit_named_asm_label_before(uid, insn);

			rtx new_set = gen_rtx_SET(dest, GEN_INT(initial_offset));
			PATTERN(insn) = new_set;
			INSN_CODE(insn) = -1; // force re-recognition
			*/

			std::cout << "Inserted labeled initial member offset " << initial_offset
				<< " for access uid " << uid << "!" << std::endl;
		}
	}
}

const pass_data unspec_to_lconst_pass_data = {
	RTL_PASS,
	"unspec_to_lconst",
	OPTGROUP_NONE,
	TV_NONE,
	PROP_rtl,
	0, 0, 0, 0
};

struct unspec_to_lconst_pass : rtl_opt_pass {
	unspec_to_lconst_pass (gcc::context *ctxt)
		: rtl_opt_pass(unspec_to_lconst_pass_data, ctxt) {}

	unsigned int execute(function* fn) override {
		unspec_to_labeled_const(fn);
		return 0;
	}
};


// Hook everything up in plugin_init

int plugin_init (struct plugin_name_args *plugin_info, struct plugin_gcc_version *version) {
	if (!plugin_default_version_check(version, &gcc_version)) {
		fprintf(stderr, "GCC version mismatch!\n");
		std::cerr << "spslr_pinpoint -> GCC version mismatch" << std::endl;
		return 1;
	}

	/*
	TODO
	vregs is almost immediately after expand (maybe the first one)
	optimizations happen afterwards (e.g. forward propagation in rtl-fwprop1)
	*/

	struct register_pass_info call_to_unspec_pass_info;
	call_to_unspec_pass_info.pass = new call_to_unspec_pass(nullptr);
	call_to_unspec_pass_info.reference_pass_name = "vregs"; // "expand";
	call_to_unspec_pass_info.ref_pass_instance_number = 1;
	call_to_unspec_pass_info.pos_op = PASS_POS_INSERT_AFTER;
	register_callback(plugin_info->base_name, PLUGIN_PASS_MANAGER_SETUP, nullptr, &call_to_unspec_pass_info);

	struct register_pass_info unspec_to_lconst_pass_info;
	unspec_to_lconst_pass_info.pass = new unspec_to_lconst_pass(nullptr);
	unspec_to_lconst_pass_info.reference_pass_name = "call_to_unspec"; // "vregs";
	unspec_to_lconst_pass_info.ref_pass_instance_number = 1;
	unspec_to_lconst_pass_info.pos_op = PASS_POS_INSERT_AFTER; // PASS_POS_INSERT_BEFORE;
	register_callback(plugin_info->base_name, PLUGIN_PASS_MANAGER_SETUP, nullptr, &unspec_to_lconst_pass_info);

	return 0;
}