#include <iostream>
#include <gcc-plugin.h>
#include <plugin-version.h>
#include <tree.h>
#include <gimple.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>

int plugin_is_GPL_compatible;

struct MemberAccessInfo {
    tree field_decl;      // FIELD_DECL des zugegriffenen Mitglieds
    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.field_decl = field;
                    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:
            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
    }
}

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

const pass_data gimplabels_pass_data = {
    GIMPLE_PASS,
    "gimplabels",
    OPTGROUP_NONE,
    TV_NONE,
    0,0,0,0,0
};

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>";


        // 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);
                for (unsigned i = 0; i < gimple_num_ops(stmt); ++i) {
                    tree op = gimple_op(stmt, i);
                    scan_tree_for_components(op, funcname, &gsi);
                }
            }
        }

		std::cout << "Number of accesses: " << g_member_accesses.size() << 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.reference_pass_name = "cfg";
    ginfo.ref_pass_instance_number = 1;
    ginfo.pos_op = PASS_POS_INSERT_AFTER;
    register_callback(plugin_info->base_name, PLUGIN_PASS_MANAGER_SETUP, nullptr, &ginfo);

    return 0;
}