#include <stdarg.h>
#include <stdio.h>
#include <stddef.h>
#include <stdint.h>

#include "common.h"

typedef struct var {
	char *name;
	intptr_t v;
} var;

#define VEC_NAME vars
#define VEC_TYPE var
#include <conts/vec.h>
typedef struct vars vars;

typedef struct proc {
	vars vars;
} proc;

static proc *gen_proc(int header, ...)
{
	(void)header;

	proc *proc = calloc(1, sizeof(struct proc));
	if (!proc)
		return NULL;

	proc->vars = vars_create(0);

	va_list args;
	va_start(args, header);
	while (1) {
		char *name = va_arg(args, char *);

		/* NULL sentinel */
		if (!name)
			break;

		var var = {
			.name = name,
			.v = 0
		};
		vars_append(&proc->vars, var);
	}
	va_end(args);

	return proc;
}

#define gen_proc(...) gen_proc(0, __VA_ARGS__, NULL)

static void proc_set(proc *proc, char *name, intptr_t v)
{
	foreach(vars, var, &proc->vars) {
		if (strcmp(var->name, name) != 0)
			continue;

		var->v = v;
		return;
	}

	/* abort? */
}

/* forward declare */
typedef struct node node;

#define VEC_NAME nodes
#define VEC_TYPE node*
#include <conts/vec.h>
typedef struct nodes nodes;

typedef struct node {
	enum kind {
		NODE_LOAD = 1,
		NODE_STORE,
		NODE_ADD,
		NODE_MUL,
		NODE_VAR,
		NODE_VARREF,
		NODE_CONST,
		NODE_ASSIGN,
		NODE_FOR,
		NODE_LT,
		NODE_BODY,
	} kind;

	/* generic values */
	struct node *n0;
	struct node *n1;
	struct node *n2;
	struct node *n3;
	intptr_t v;

	/* list of nodes. If we really tried to save space, we could make all
	 * nodes just be indexes into this nodes array, but that would mean that
	 * execution has to go through more layers of indirection, so I'm
	 * reserving n0/n1/n2/n3 for quicker lookup */
	nodes nodes;

	/* used by var/varref nodes to point to variable in proc context */
	intptr_t *varref;
} node;

static node *gen_node(enum kind kind, node *n0, node *n1, node *n2, node *n3, intptr_t v)
{
	node *n = calloc(1, sizeof(node));
	if (!n)
		return NULL;

	/* unused by generic nodes */
	n->nodes = nodes_create(0);

	n->kind = kind;
	n->n0 = n0;
	n->n1 = n1;
	n->n2 = n2;
	n->n3 = n3;
	n->v = v;
	return n;
}


#define get_n0(k, x) ({assert((x)->kind == (k)); (x)->n0;})
#define get_n1(k, x) ({assert((x)->kind == (k)); (x)->n1;})
#define get_n2(k, x) ({assert((x)->kind == (k)); (x)->n2;})
#define get_n3(k, x) ({assert((x)->kind == (k)); (x)->n3;})
#define get_v(k, x) ({assert((x)->kind == (k)); (x)->v;})

/* generic defines for 'trivial' nodes */

#define gen_load(addr) gen_node(NODE_LOAD, (addr), NULL, NULL, NULL, 0)
#define load_addr(x) get_n0(NODE_LOAD, (x))

#define gen_store(addr, val) gen_node(NODE_STORE, (addr), (val), NULL, NULL, 0)
#define store_addr(x) get_n0(NODE_STORE, (x))
#define store_val(x) get_n1(NODE_STORE, (x))

#define gen_add(l, r) gen_node(NODE_ADD, (l), (r), NULL, NULL, 0)
#define add_l(x) get_n0(NODE_ADD, (x))
#define add_r(x) get_n1(NODE_ADD, (x))

#define gen_mul(l, r) gen_node(NODE_MUL, (l), (r), NULL, NULL, 0)
#define mul_l(x) get_n0(NODE_MUL, (x))
#define mul_r(x) get_n1(NODE_MUL, (x))

#define gen_const(c) gen_node(NODE_CONST, NULL, NULL, NULL, NULL, (c))
#define const_v(x) get_v(NODE_CONST, (x))

#define gen_assign(dst, src) gen_node(NODE_ASSIGN, (dst), (src), NULL, NULL, 0)
#define assign_dst(x) get_n0(NODE_ASSIGN, (x))
#define assign_src(x) get_n1(NODE_ASSIGN, (x))

#define gen_for(init, cond, post, body) gen_node(NODE_FOR, (init), (cond), (post), (body), 0)
#define for_init(x) get_n0(NODE_FOR, (x))
#define for_cond(x) get_n1(NODE_FOR, (x))
#define for_post(x) get_n2(NODE_FOR, (x))
#define for_body(x) get_n3(NODE_FOR, (x))

#define gen_lt(l, r) gen_node(NODE_LT, (l), (r), NULL, NULL, 0)
#define lt_l(x) get_n0(NODE_LT, (x))
#define lt_r(x) get_n1(NODE_LT, (x))

/* nodes that need special handling */
static node *gen_body(int header, ...)
{
	node *body = calloc(1, sizeof(node));
	if (!body)
		return NULL;

	body->kind = NODE_BODY;
	body->nodes = nodes_create(0);

	va_list args;
	va_start(args, header);
	while (1) {
		node *n = va_arg(args, node *);
		if (!n)
			break;

		nodes_append(&body->nodes, n);
	}
	va_end(args);

	return body;
}

#define gen_body(...) gen_body(0, __VA_ARGS__, NULL)
#define body_list(x) ({assert((x)->kind == NODE_BODY); (x)->nodes;})

static node *gen_var(proc *proc, char *name)
{
	assert(proc && name);

	node *n = calloc(1, sizeof(node));
	if (!n)
		return NULL;

	n->kind = NODE_VAR;
	foreach(vars, var, &proc->vars) {
		if (strcmp(var->name, name) != 0)
			continue;

		/* assumes that the var list is static, which it really should
		 * be by this point */
		n->varref = &var->v;
		return n;
	}

	/* no such var (abort or return some more obvious error) */
	free(n);
	return NULL;
}

#define var(x) ({assert((x)->kind == NODE_VAR); (x)->varref;})

static node *gen_varref(proc *proc, char *name)
{
	/* node generation logic is the same, it's just the runtime that's
	 * different from gen_var, so just change NODE_VAR -> NODE_VARREF */
	node *n = gen_var(proc, name);
	if (!n)
		return NULL;

	n->kind = NODE_VARREF;
	return n;
}

#define varref(x) ({assert((x)->kind == NODE_VARREF); (x)->varref;})

/* predeclare */
static intptr_t run(node *n);

static intptr_t run_load(node *n)
{
	return *(intptr_t* )run(load_addr(n));
}

static intptr_t run_store(node *n)
{
	intptr_t *dst = (intptr_t *)run(store_addr(n));
	intptr_t val = run(store_val(n));
	*dst = val;
	return 0;
}

static intptr_t run_add(node *n)
{
	intptr_t l = run(add_l(n));
	intptr_t r = run(add_r(n));
	return l + r;
}

static intptr_t run_mul(node *n)
{
	intptr_t l = run(mul_l(n));
	intptr_t r = run(mul_r(n));
	return l * r;
}

static intptr_t run_var(node *n)
{
	return *var(n);
}

static intptr_t run_varref(node *n)
{
	return (intptr_t)varref(n);
}

static intptr_t run_const(node *n)
{
	return const_v(n);
}

static intptr_t run_assign(node *n)
{
	intptr_t *dst = (intptr_t *)run(assign_dst(n));
	intptr_t src = run(assign_src(n));
	*dst = src;
	return 0;
}

static intptr_t run_for(node *n)
{
	for (run(for_init(n)); run(for_cond(n)); run(for_post(n)))
		run(for_body(n));

	return 0;
}

static intptr_t run_lt(node *n)
{
	intptr_t l = run(lt_l(n));
	intptr_t r = run(lt_r(n));
	return l < r;
}

static intptr_t run_body(node *n)
{
	nodes list = body_list(n);
	foreach(nodes, node, &list) {
		run(*node);
	}

	return 0;
}

static intptr_t run(node *n)
{
	switch (n->kind) {
	case NODE_LOAD:   return run_load(n);
	case NODE_STORE:  return run_store(n);
	case NODE_ADD:    return run_add(n);
	case NODE_MUL:    return run_mul(n);
	case NODE_VAR:    return run_var(n);
	case NODE_VARREF: return run_varref(n);
	case NODE_CONST:  return run_const(n);
	case NODE_ASSIGN: return run_assign(n);
	case NODE_FOR:    return run_for(n);
	case NODE_LT:     return run_lt(n);
	case NODE_BODY:   return run_body(n);
	default: fprintf(stderr, "unimp\n");
		 abort();
	}

	return 0;
}