#include <assert.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>

#include <linux/ptrace.h>

#include "arch.h"
#include "func.h"
#include "ir.h"
#include "node.h"
#include "ply.h"
#include "sym.h"
#include "type.h"

struct kprobe {
};

/* regs */

static struct type t_pt_regsp = {
	.ttype = T_POINTER,

	.ptr.type = &t_pt_regs,
};

static const struct func kprobe_regs_func = {
	.name = "regs",

	.type = &t_pt_regsp,
	.static_ret = 1,
};


/* argN */

static inline int is_arg(const char *name)
{
	return (strstr(name, "arg") == name)
		&& (strlen(name) == 4)
		&& (name[3] >= '0' && name[3] <= '9');
}

static int kprobe_arg_rewrite(const struct func *func, struct node *n,
			      struct prog *prog)
{
	struct node *new;
	const char *reg;
	int arg;

	arg = n->ident.name[3] - '0';
	reg = arch_register_argument(arg);
	if (!reg) {
		_e("%#N: the location of %N is unknown\n", n, n);

		/* TODO: add ABI mappings for specifying arguments
		 * passed on the stack. */
		return -EINVAL;
	}

	/* argN => (*regs).REG */
	new = node_expr(".",
			node_expr(":deref", node_ident("regs"), NULL),
			node_string((char *)reg),
			NULL);

	node_replace(n, new);
	return 0;
}

static const struct func kprobe_arg_func = {
	.name = "argN",

	/* for now, in the future we could read dwarf symbols to
	 * figure out the real type. */
	.type = &t_ulong,
	.static_ret = 1,

	.rewrite = kprobe_arg_rewrite,
};


/*  */

static int kprobe_ir_pre(struct prog *prog)
{
	struct sym **sym;

	symtab_foreach(prog->locals, sym) {
		if ((*sym)->name && !strcmp((*sym)->name, "regs")) {
			ir_init_sym(prog->ir, *sym);

			/* kernel sets r1 to the address of the
			 * pt_regs struct, which ply denotes as
			 * 'regs'. if we're using it we need to get a
			 * reference to it before it is clobbered. */
			ir_emit_insn(prog->ir, MOV(0, 0),
				     (*sym)->irs.reg, BPF_REG_1);
		}
	}

	return 0;
}

static int kprobe_sym_alloc(struct prog *prog, struct node *n)
{
	const struct func *func = NULL;
	int err;

	switch (n->ntype) {
	case N_EXPR:
		break;
	case N_IDENT:
		if (!strcmp(n->ident.name, "regs"))
			func = &kprobe_regs_func;
		else if (is_arg(n->ident.name))
			func = &kprobe_arg_func;
		break;
	default:
		break;
	}

	if (!func)
		return -ENOENT;

	err = func_static_validate(func, n);
	if (err)
		return err;

	n->sym = sym_alloc(prog->locals, n, func);

	if (func->static_ret)
		n->sym->type = func_return_type(func);
	return 0;
}

static int kprobe_probe(struct prog *prog)
{
	struct kprobe *kp;

	kp = calloc(1, sizeof(*kp));
	assert(kp);

	prog->provider_data = kp;
	return 0;
}

struct provider kprobe = {
	.name = "kprobe",

	.ir_pre = kprobe_ir_pre,
	/* .rewrite_node = kprobe_rewrite_node, */
	/* .resolve = kprobe_resolve, */
	.sym_alloc = kprobe_sym_alloc,
	.probe = kprobe_probe,
};

__attribute__((constructor))
static void kprobe_init(void)
{
	provider_register(&kprobe);
}