#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 <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>
#include <unordered_map>

int plugin_is_GPL_compatible;

struct MemberOffset {
    using UID = unsigned long;

    std::string type;
    std::string member;
    unsigned long offset;
};

static MemberOffset::UID next_moffset_uid = 0;
static std::unordered_map<MemberOffset::UID, MemberOffset> member_offset_occasions;

static void print_gimple_tree(tree t, int depth, int indent = 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;

    for (int i = 0; i < TREE_CODE_LENGTH(TREE_CODE(t)); ++i)
        print_gimple_tree(TREE_OPERAND(t, i), depth + indent, indent);
}

static void print_gimple_statement(gimple* stmt) {
    if (!stmt)
        return;

    print_gimple_stmt(stdout, stmt, 0, TDF_NONE);

    for (unsigned i = 0; i < gimple_num_ops(stmt); ++i) {
        tree operand = gimple_op(stmt, i);
        print_gimple_tree(operand, 2, 2);
    }
}

static bool is_member_offset(gimple* stmt, MemberOffset::UID& uid) {
	if (!is_gimple_call(stmt))
		return false;

	tree fndecl = gimple_call_fndecl(stmt);
	if (!fndecl || !DECL_NAME(fndecl))
		return false;

	const char *fname = IDENTIFIER_POINTER(DECL_NAME(fndecl));
	if (!fname || strcmp(fname, "__get_member_offset") != 0)
		return false;

	if (gimple_call_num_args(stmt) != 1)
		return false;

	tree farg0 = gimple_call_arg(stmt, 0);
	if (!farg0 || TREE_CODE(farg0) != INTEGER_CST)
		return false;

	uid = (MemberOffset::UID) tree_to_uhwi(farg0);
	return true;
}

static int scan_gimple_statement(const char* funcname, gimple_stmt_iterator* gsi) {
    print_gimple_statement(gsi_stmt(*gsi));

	// Replace calls (later do it in RTL with UNSPECs)
	MemberOffset::UID moffset_uid;
	if (is_member_offset(gsi_stmt(*gsi), moffset_uid)) {
		auto moit = member_offset_occasions.find(moffset_uid);
		if (moit == member_offset_occasions.end()) {
			std::cout << "Failed to find member offset occasion data!" << std::endl;
			return 1;
		}

		const MemberOffset& mo = moit->second;

		tree lhs = gimple_get_lhs(gsi_stmt(*gsi));
		tree ctype = lhs ? TREE_TYPE(lhs) : size_type_node;

		tree cst = build_int_cst(ctype, mo.offset);

		if (lhs) {
			gimple *assign = gimple_build_assign(lhs, cst);
			gsi_replace(gsi, assign, true);
			std::cout << "REPLACED MOFF " << moffset_uid << " -> offsetof(" << mo.type << ", " << mo.member << ")" << std::endl;
			print_gimple_statement(gsi_stmt(*gsi));
		} else {
			gsi_remove(gsi, true);
			std::cout << "REMOVED MOFF " << moffset_uid << " -> offsetof(" << mo.type << ", " << mo.member << ")" << std::endl;
		}
	}

    return 0;
}

static const pass_data access_discover_pass_data = {
    GIMPLE_PASS,
    "access_discover",
    OPTGROUP_NONE,
    TV_NONE,
    0,0,0,0,
    TODO_update_ssa | TODO_cleanup_cfg | TODO_verify_il
};

struct access_discover_pass : gimple_opt_pass {
    access_discover_pass(gcc::context *ctxt);
    unsigned int execute(function* fun) override;
};

access_discover_pass::access_discover_pass(gcc::context *ctxt) : gimple_opt_pass(access_discover_pass_data, ctxt) {}

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

    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)) {
            if (scan_gimple_statement(funcname, &gsi) != 0) {
                internal_error("fatal error in spslr plugin: %s", "failed to scan gimple statement");
                return 0;
            }
        }
    }

    return 0;
}

static tree my_component_ref_cb (tree ref)
{
  tree base  = TREE_OPERAND (ref, 0);   /* p           */
  tree field = TREE_OPERAND (ref, 1);   /* FIELD_DECL  */
  tree basetype = TREE_TYPE (base);     /* struct type */

  if (!basetype || TREE_CODE (basetype) != RECORD_TYPE)
    return NULL_TREE;  /* Not a struct, leave alone.  */

  /* Choose which structs you care about.  */
  const char *name = NULL;
  if (TYPE_NAME (basetype))
    {
      tree tn = TYPE_NAME (basetype);
      if (TREE_CODE (tn) == TYPE_DECL && DECL_NAME (tn))
        name = IDENTIFIER_POINTER (DECL_NAME (tn));
      else if (TREE_CODE (tn) == IDENTIFIER_NODE)
        name = IDENTIFIER_POINTER (tn);
    }

  if (!name || strcmp (name, "task_struct") != 0)
    return NULL_TREE; /* Only modify struct Tracked.  */

  location_t loc = EXPR_LOCATION (ref);
  tree result_type = TREE_TYPE (ref);        /* type of the field access */

  /* 1. Build the __get_member_offset call. TODO -> Do not redo it every time...  */
  tree param_types =
      tree_cons (NULL_TREE, long_unsigned_type_node, NULL_TREE);
  tree fntype = build_function_type (sizetype, param_types);

  tree fn = build_fn_decl ("__get_member_offset", fntype);
  DECL_EXTERNAL (fn) = 1;
  TREE_PUBLIC (fn) = 1;
  DECL_ARTIFICIAL (fn) = 1;

  // Prevent VOP problems later when removing calls (VOPs manage memory effects, which these calls have none of)
  DECL_PURE_P(fn) = 1;
  DECL_IS_NOVOPS(fn) = 1;

  tree off_tree = byte_position(field);
  if (!off_tree || TREE_CODE(off_tree) != INTEGER_CST) {
    std::cout << "Failed to get member offset!" << std::endl;
    return NULL_TREE;
  }

  unsigned long member_offset = (unsigned long) tree_to_uhwi(off_tree);

  std::string member_name;

  tree member_name_id = DECL_NAME(field);
  if (!member_name_id)
    member_name = std::string{ "<anonymous>" };  // e.g. unnamed bitfields

  member_name = std::string{ IDENTIFIER_POINTER(member_name_id) };

  unsigned long next_uid = next_moffset_uid++;

  MemberOffset moffset;
  moffset.type = std::string{ name };
  moffset.member = member_name;
  moffset.offset = member_offset;

  member_offset_occasions.emplace(next_uid, moffset);

  tree moffset_uid = build_int_cst (long_unsigned_type_node, next_uid);

  tree offset_call = build_call_expr_loc (loc, fn, 1, moffset_uid);
  // TREE_SIDE_EFFECTS (offset_call) = 1; // Maybe not

  /* 2. Get a (char *) pointer to the base.  */
  tree char_ptr_type = build_pointer_type (char_type_node);
  tree base_ptr = NULL_TREE;

  if (POINTER_TYPE_P (TREE_TYPE (base)))
    {
	  // TODO -> Is it actually a pointer? INDIRECT_REF is used, so it is dereferenced already!
      /* Case: p->member, base is already a pointer.  */
      base_ptr = build1_loc (loc, NOP_EXPR, char_ptr_type,
                             build1_loc (loc, NOP_EXPR,
                                         TREE_TYPE (char_ptr_type), base));
    }
  else
    {
      /* Case: T.member, take its address.  */
      tree base_addr = build1_loc (loc, ADDR_EXPR,
                                   build_pointer_type (TREE_TYPE (base)),
                                   base);
	  // TODO -> case of INDIRECT_REF is 2 blocks that cancel each other out -> &*ptr
      base_ptr = build1_loc (loc, NOP_EXPR, char_ptr_type, base_addr);
    }

  /* 3. Add offset and cast back to field pointer, then dereference.  */
  tree plus = build2_loc (loc, POINTER_PLUS_EXPR, char_ptr_type,
                          base_ptr, offset_call);

  tree field_ptr_type = build_pointer_type (result_type);
  tree cast_back = build1_loc (loc, NOP_EXPR, field_ptr_type, plus);

  tree new_ref = build1_loc (loc, INDIRECT_REF, result_type, cast_back);

  /* Optional: mark to suppress further folding.  */
  // TREE_SIDE_EFFECTS (new_ref) = 1;

  return new_ref; // ALWAYS an INDIRECT_REF, directly to the member (same type as COMPONENT_REF!)
}

static void build_component_ref_cb(void* event_data, void* user_data) {
	tree* component_ref_node = (tree*)event_data;
	if (!component_ref_node)
		return;

	tree repl = my_component_ref_cb(*component_ref_node);
	if (repl)
		*component_ref_node = repl;
}

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

	register_callback(plugin_info->base_name, PLUGIN_BUILD_COMPONENT_REF, build_component_ref_cb, NULL);

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

	return 0;
}