#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 {
	int use_ctx:1;

	struct irstate ctx_irs;
};

/* ctx */

static struct type t_pt_regsp = {
	.ttype = T_POINTER,

	.ptr.type = &t_pt_regs,
};

static const struct func kprobe_ctx_func = {
	.name = "ctx",

	.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 => (*ctx).REG */
	new = node_expr(".",
			node_expr("*", node_ident("ctx"), NULL),
			node_string((char *)reg),
			NULL);

	ast_fprint(stdout, n);
	ast_fprint(stdout, new);

	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 kprobe *kp = prog->provider_data;

	if (!kp->use_ctx)
		return 0;

	irs_alloc_reg(&kp->ctx_irs, prog->ir);

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

static int kprobe_sym_alloc(struct prog *prog, struct node *n)
{
	struct kprobe *kp = prog->provider_data;
	const struct func *func = NULL;
	int err;

	switch (n->ntype) {
	case N_EXPR:
		break;
	case N_IDENT:
		if (!is_arg(n->ident.name))
			break;

		func = &kprobe_arg_func;
		kp->use_ctx = 1;
		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);
}