#include <iostream>
#include <gcc-plugin.h>
#include <plugin-version.h>
#include <tree.h>
#include <gimple.h>
#include <gimplify.h>
#include <gimple-iterator.h>
#include <basic-block.h>
#include <gimple-ssa.h>
#include <tree-pass.h>
#include <cctype>
#include <cstdio>
#include <vector>
#include <string>
#include <unordered_set>
#include <unordered_map>
#include <pretty-print.h>
#include <tree-pretty-print.h>
#include <diagnostic.h>
#include <gimple-pretty-print.h>
#include <obstack.h>
#include <cgraph.h>
#include <context.h>
#include <function.h>
#include <builtins.h>

int plugin_is_GPL_compatible;

struct MemberAccessInfo {
    const char* sname;
	unsigned long label_num;
    std::string funcname;
    location_t loc;
};

static std::vector<MemberAccessInfo> g_member_accesses;
static unsigned long label_counter = 0;
std::unordered_set<std::string> targets;

bool is_target(const char* name) {
	std::string str{ name };
	return targets.find(str) != targets.end();
}

// ------------------------
// Sicherer Helfer zum Abrufen des RECORD_TYPE-Namens
// ------------------------
static const char* safe_get_record_type_name_from_field(tree field)
{
    if (!field) return nullptr;

    tree ctx = DECL_CONTEXT(field);
    int depth = 0;
    while (ctx && depth < 8 && TREE_CODE(ctx) != RECORD_TYPE) {
        if (TREE_CODE(ctx) == TYPE_DECL)
            ctx = TREE_TYPE(ctx);
        else
            ctx = nullptr;
        depth++;
    }

    if (!ctx || TREE_CODE(ctx) != RECORD_TYPE)
        return nullptr;

    tree type_name = TYPE_NAME(ctx);
    if (!type_name)
        return nullptr;

    tree decl_name = nullptr;
    if (TREE_CODE(type_name) == TYPE_DECL)
        decl_name = DECL_NAME(type_name);
    else if (TREE_CODE(type_name) == IDENTIFIER_NODE)
        decl_name = type_name;

    if (!decl_name || TREE_CODE(decl_name) != IDENTIFIER_NODE)
        return nullptr;

    return IDENTIFIER_POINTER(decl_name);
}

// ------------------------
// Rekursiver Scanner für COMPONENT_REF
// ------------------------
static void scan_tree_for_components(tree t, const char *funcname, gimple_stmt_iterator *gsi)
{
    if (!t) return;

    tree_code tc = TREE_CODE(t);

    switch(tc) {
        case COMPONENT_REF: {
            tree field = TREE_OPERAND(t,1);
            if (field && TREE_CODE(field) == FIELD_DECL) {
                const char* rec_name = safe_get_record_type_name_from_field(field); // struct name
                if (rec_name && is_target(rec_name)) {
                    MemberAccessInfo info;
                    info.sname = rec_name;
                    info.label_num = ++label_counter;
                    info.funcname = funcname ? funcname : "<unknown>";
                    info.loc = UNKNOWN_LOCATION;
                    g_member_accesses.push_back(info);

                    // insert label immediately
                    char buf[64];
                    snprintf(buf, sizeof(buf), "mylabel_%lu:", info.label_num);

                    vec<tree, va_gc>* inputs   = nullptr;
                    vec<tree, va_gc>* outputs  = nullptr;
                    vec<tree, va_gc>* clobbers = nullptr;
                    vec<tree, va_gc>* labels   = nullptr;

                    gimple *asm_stmt = gimple_build_asm_vec(
                        ggc_strdup(buf),
                        inputs,
                        outputs,
                        clobbers,
                        labels
                    );

                    gsi_insert_before(gsi, asm_stmt, GSI_SAME_STMT);
                }
            }
            // Rekursion in das Basisobjekt
            scan_tree_for_components(TREE_OPERAND(t,0), funcname, gsi);
            break;
        }
        case INDIRECT_REF:
        case CONVERT_EXPR:
        case NOP_EXPR:
        case VIEW_CONVERT_EXPR:
		case ADDR_EXPR:
            scan_tree_for_components(TREE_OPERAND(t,0), funcname, gsi);
            break;
        case ARRAY_REF:
            scan_tree_for_components(TREE_OPERAND(t,0), funcname, gsi);
            scan_tree_for_components(TREE_OPERAND(t,1), funcname, gsi);
            break;
        case PLUS_EXPR:
        case MINUS_EXPR:
        case MULT_EXPR:
        case POINTER_PLUS_EXPR:
            scan_tree_for_components(TREE_OPERAND(t,0), funcname, gsi);
            scan_tree_for_components(TREE_OPERAND(t,1), funcname, gsi);
            break;
        default:
            break; // Blätter
    }
}

void scan_stmt_for_offsetof(const char *funcname, gimple_stmt_iterator *gsi) {
	gimple *stmt = gsi_stmt(*gsi);
	if (!is_gimple_call(stmt))
		return;

	tree callee = gimple_call_fn(stmt);
	if (!callee)
		return;

	if (TREE_CODE(callee) == ADDR_EXPR)
		callee = TREE_OPERAND(callee, 0);

	if (!callee || TREE_CODE(callee) != FUNCTION_DECL)
		return;

	const char* name = IDENTIFIER_POINTER(DECL_NAME(callee));
	if (!name || strcmp(name, "__spslr_offsetof") != 0)
		return;

	tree arg_type   = gimple_call_arg(stmt, 0);  // struct X
	tree arg_member = gimple_call_arg(stmt, 1);  // member
	tree arg_value  = gimple_call_arg(stmt, 2);  // (likely folded) offsetof value

	if (TREE_CODE(arg_type) == ADDR_EXPR)
		arg_type = TREE_OPERAND(arg_type, 0);

	if (TREE_CODE(arg_member) == ADDR_EXPR)
		arg_member = TREE_OPERAND(arg_member, 0);

	if (TREE_CODE(arg_type) != STRING_CST || TREE_CODE(arg_member) != STRING_CST) {
		std::cout << "Failed to parse __spslr_offsetof arguments!" << std::endl;
		return;
	}

	const char* type_str = TREE_STRING_POINTER(arg_type); // includes "struct "
	const char* member_str = TREE_STRING_POINTER(arg_member);

	// Remove internal call
	tree lhs = gimple_call_lhs(stmt);
	tree val = gimple_call_arg(stmt, 2);

	if (!lhs) {
		gsi_remove(gsi, true); // returns void -> just remove call entirely
		return;
	}

	val = fold_convert(TREE_TYPE(lhs), val);

	gassign *as = gimple_build_assign(lhs, val);
	gsi_replace(gsi, as, true);

    if (strncmp(type_str, "struct ", 7) == 0)
        type_str = type_str + 7;

	if (!is_target(type_str))
        return;

	// Insert label
	MemberAccessInfo info;
	info.sname = type_str;
	info.label_num = ++label_counter;
	info.funcname = funcname ? funcname : "<unknown>";
	info.loc = UNKNOWN_LOCATION;
	g_member_accesses.push_back(info);

	char buf[64];
	snprintf(buf, sizeof(buf), "mylabel_%lu:", info.label_num);

	vec<tree, va_gc>* inputs   = nullptr;
	vec<tree, va_gc>* outputs  = nullptr;
	vec<tree, va_gc>* clobbers = nullptr;
	vec<tree, va_gc>* labels   = nullptr;

	gimple *asm_stmt = gimple_build_asm_vec(
		ggc_strdup(buf),
		inputs,
		outputs,
		clobbers,
		labels
	);

	gsi_insert_before(gsi, asm_stmt, GSI_SAME_STMT);
}

// ------------------------
// GIMPLE-Pass-Ausführung
// ------------------------
namespace {

const pass_data gimplabels_pass_data = {
    GIMPLE_PASS,
    "gimplabels",
    OPTGROUP_NONE,
    TV_NONE,
    0,0,0,0,
	TODO_update_ssa | TODO_cleanup_cfg | TODO_verify_il
};

static void
plugin_debug_tree(tree t, int depth = 2)
{
    if (!t)
        return;

    for (int i = 0; i < depth; ++i)
        std::cout << " ";

    std::cout << get_tree_code_name(TREE_CODE(t));

    if (TREE_CODE(t) == FIELD_DECL && DECL_NAME(t))
        std::cout << " <" << IDENTIFIER_POINTER(DECL_NAME(t)) << ">";
    else if (TREE_CODE(t) == SSA_NAME && SSA_NAME_VAR(t)
             && DECL_NAME(SSA_NAME_VAR(t)))
        std::cout << " <ssa " << IDENTIFIER_POINTER(DECL_NAME(SSA_NAME_VAR(t))) << ">";

    std::cout << std::endl;

    /* Recurse over child operands */
    for (int i = 0; i < TREE_CODE_LENGTH(TREE_CODE(t)); ++i)
        plugin_debug_tree(TREE_OPERAND(t, i), depth + 2);
}

struct gimplabels_pass : gimple_opt_pass {
    gimplabels_pass(gcc::context *ctxt)
        : gimple_opt_pass(gimplabels_pass_data, ctxt)
    {}

    unsigned int execute(function* fun) override {
        const char* funcname = fun->decl && DECL_NAME(fun->decl)
                               ? IDENTIFIER_POINTER(DECL_NAME(fun->decl))
                               : "<anon>";

        constexpr bool verbose = false;
        std::size_t base_access_count = g_member_accesses.size();
        std::cout << (verbose ? "\n" : "") << "Handling function \"" << funcname << "\"..." << std::endl;

        // Iteration über alle Grundblöcke und Anweisungen
        basic_block bb;
        FOR_EACH_BB_FN(bb, fun) {
            for (gimple_stmt_iterator gsi = gsi_start_bb(bb); !gsi_end_p(gsi); gsi_next(&gsi)) {
                gimple *stmt = gsi_stmt(gsi);
                if (verbose)
                    print_gimple_stmt (stdout, stmt, 0, TDF_NONE);
                scan_stmt_for_offsetof(funcname, &gsi);
                for (unsigned i = 0; i < gimple_num_ops(stmt); ++i) {
                    tree op = gimple_op(stmt, i);
                    if (verbose)
                        plugin_debug_tree(op);
                    scan_tree_for_components(op, funcname, &gsi);
                }
            }
        }

        std::cout << "    Number of accesses: " << (g_member_accesses.size() - base_access_count) << std::endl;
        return 0;
    }
};

} // anonymer Namensraum

// ------------------------
// Plugin-Initialisierung
// ------------------------

static tree handle_slr_attr(tree* node, tree name, tree args, int flags, bool* no_add_attrs) {
    if (!node)
        return NULL_TREE;

    if (TREE_CODE(*node) != RECORD_TYPE)
        return *node; // Only allowed on structs

    tree type_name_tree = TYPE_IDENTIFIER(*node);
    if (type_name_tree == NULL_TREE)
        return *node;

    std::string name_str{ IDENTIFIER_POINTER(type_name_tree) };
    targets.insert(name_str);

	std::cout << "SLR attribute recognized on \"" << name_str << "\"" << std::endl;
    return *node;
}

static struct attribute_spec slr_attr = { "slr", 0, 0, false, false, false, false, handle_slr_attr, NULL };

void register_attributes(void* event_data, void* data) {
    register_attribute(&slr_attr);
}

int plugin_init(plugin_name_args *plugin_info, plugin_gcc_version *version)
{
    if (!plugin_default_version_check(version, &gcc_version)) {
        fprintf(stderr, "GCC-Version stimmt nicht überein\n");
        return 1;
    }

	register_callback(plugin_info->base_name, PLUGIN_ATTRIBUTES, register_attributes, NULL);

    g_member_accesses.clear();
    label_counter = 0;

    struct register_pass_info ginfo;
    ginfo.pass = new gimplabels_pass(nullptr);
    ginfo.ref_pass_instance_number = 1;
    ginfo.reference_pass_name = "optimized";
    ginfo.pos_op = PASS_POS_INSERT_AFTER;
    register_callback(plugin_info->base_name, PLUGIN_PASS_MANAGER_SETUP, nullptr, &ginfo);

    return 0;
}