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

#include <linux/ptrace.h>

#include "ply.h"

struct kprobe {
};

static int kprobe_ir_prologue(prog_t *prog)
{
	sym_t *ctx = sym_get(prog->locals, "ctx");

	if (!ctx)
		return 0;

	irs_alloc_reg(&ctx->irs, prog->ir);

	/* kernel sets r1 to the address of the context */
	ir_emit_insn(prog->ir, MOV(0, 0), ctx->irs.reg, BPF_REG_1);
	return 0;
}

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_rewrite_arg(prog_t *prog, node_t *n)
{
	const char *reg;
	int arg = n->ident[3] - '0';
	node_t *new, *ctx;

	reg = arch_register_argument(arg);
	if (!reg) {
		node_print(n, stderr);
		fputs(": the location of this argument is unknown\n", stderr);
		/* TODO: add ABI mappings for specifying arguments
		 * passed on the stack. */
		return -EINVAL;
	}

	ctx = node_ident("ctx");

	/* argN => (*ctx).REG */
	new = node_vlist(node_keyword('.', 0),
			 node_vlist(node_keyword('*', 0),
				    ctx,
				    NULL),
			 node_string(reg),
			 NULL);

	ctx->type = type_ptr_of(&t_pt_regs);
	new->type = n->type;
	new->list->type = &t_void;
	new->list->next->type = &t_pt_regs;
	new->list->next->list->type = &t_void;
	node_replace(n, new);

	return sym_add(prog->locals, ctx->ident, ctx->type, &ctx->sym);
}

static int kprobe_rewrite_node(prog_t *prog, node_t *n)
{
	if ((n->ntype == N_IDENT) && is_arg(n->ident))
		return kprobe_rewrite_arg(prog, n);

	return 0;
}

static int kprobe_resolve(prog_t *prog, node_t *n)
{
	type_t *t;

	if (is_arg(n->ident))
		t = &t_ulong;
	else if (!strcmp(n->ident, "ctx"))
		t = type_ptr_of(&t_pt_regs);
	else
		return -ENOENT;

	n->type = t;
	return sym_add(prog->locals, n->ident, t, &n->sym);
}

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

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

	prog->provider_data = kp;
	return 0;
}

provider_t kprobe = {
	.name = "kprobe",

	.ir_prologue = kprobe_ir_prologue,
	.rewrite_node = kprobe_rewrite_node,
	.resolve = kprobe_resolve,
	.probe = kprobe_probe,
};

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