#include "ast.h"
#include "slice.h"
#include "vector.h"

#include <stdio.h>
#include <string.h>

#define MEMCPY(d, s) memcpy(d, s, sizeof(*d))

#define MEMDUP(T, d, s) T d = malloc(sizeof(*d)); memcpy(d, s, sizeof(*d))

struct slice 
expression_token_literal(struct expression *expr)
{
	return expr->token.literal;
}

struct slice 
statement_token_literal(struct statement *st)
{
	return st->token.literal;
}

struct slice 
program_token_literal(struct program *prog)
{
	if (prog->statements->len > 0) {
		return node_token_literal(
				(struct statement *)prog->statements->values[0]);
	} else {
		return slice_new("", 0, 0);
	}
}

static inline char *
let_statement_sprint(struct let_statement *st, char *str)
{
	char litbuf[32];
	char namebuf[32];
	sprintf(str, "%s %s = ", slice_sprint(&st->token.literal, litbuf),
			slice_sprint(&st->name->value, namebuf));
	char *cur = str + strlen(str);
	expression_sprint(st->value, cur);
	cur += strlen(cur);
	cur[0] = ';', cur[1] = '\0';
	return str;
}

static inline char *
return_statement_sprint(struct return_statement *st, char *str)
{
	char litbuf[32];
	sprintf(str, "%s ", slice_sprint(&st->token.literal, litbuf));
	char *cur = str + strlen(str);
	expression_sprint(st->value, cur);
	cur += strlen(cur);
	cur[0] = ';', cur[1] = '\0';
	return str;
}

static inline char *
block_statement_sprint(struct block_statement *st, char *str)
{
	size_t i;
	struct statement *v;
	char *cur = str;
	vector_foreach(st->statements, i, v) {
		statement_sprint(v, cur);
		cur += strlen(cur);
	}
	return str;
}

static inline char *
expression_statement_sprint(struct expression_statement *st, char *str)
{
	return expression_sprint(st->expr, str);
}

static inline char *
prefix_expression_sprint(struct prefix_expression *expr, char *str)
{
	*str = '(';
	char *cur = str + 1;
	slice_sprint(&expr->operator, cur);
	cur += strlen(cur);
	expression_sprint(expr->right, cur);
	cur += strlen(cur);
	cur[0] = ')', cur[1] = '\0';
	return str;
}

static inline char *
infix_expression_sprint(struct infix_expression *expr, char *str)
{
	*str = '(';
	char *cur = str + 1;
	expression_sprint(expr->left, cur);
	cur += strlen(cur);
	*cur = ' ', cur++;
	slice_sprint(&expr->operator, cur);
	cur += strlen(cur);
	*cur = ' ', cur++;
	expression_sprint(expr->right, cur);
	cur += strlen(cur);
	cur[0] = ')', cur[1] = '\0';
	return str;
}

static inline char *
if_expression_sprint(struct if_expression *expr, char *str)
{
	char *cur = str;
	strcpy(cur, "if ");
	cur += strlen(cur);
	expression_sprint(expr->condition, cur);
	cur += strlen(cur);
	*cur = ' ', cur++;
	statement_sprint(expr->consequence, cur);
	if (expr->alternative) {
		cur += strlen(cur);
		strcpy(cur, "else ");
		cur += strlen(cur);
		statement_sprint(expr->alternative, cur);
	}

	return str;
}

static inline char *
func_literal_sprint(struct func_literal *expr, char *str)
{
	char *cur = str;
	slice_sprint(&expr->token.literal, cur);
	cur += strlen(cur);
	*cur = '(', cur++;
	size_t i;
	struct expression *ident;
	vector_foreach(expr->parameters, i, ident) {
		slice_sprint(&ident->token.literal, cur);
		cur += strlen(cur);
		if (i < (expr->parameters->len - 1)) {
			cur[0] = ',', cur[1] = ' ', cur += 2;
		}
	}
	*cur = ')', cur++;
	statement_sprint(expr->body, cur);

	return str;
}

static inline char *
call_expression_sprint(struct call_expression *expr, char *str)
{
	char *cur = str;
	expression_sprint(expr->func, cur);
	cur += strlen(cur);
	*cur = '(', cur++;
	size_t i;
	struct expression *arg;
	vector_foreach(expr->arguments, i, arg) {
		expression_sprint(arg, cur);
		cur += strlen(cur);
		if (i < (expr->arguments->len - 1)) {
			cur[0] = ',', cur[1] = ' ', cur += 2;
		}
	}
	cur[0] = ')', cur[1] = '\0';

	return str;
}

char *
expression_sprint(struct expression *expr, char *str)
{
	switch (expr->type) {
		case EXPRESSION_PREFIX:
			return prefix_expression_sprint(&expr->prefix, str);
		case EXPRESSION_INFIX:
			return infix_expression_sprint(&expr->infix, str);
		case EXPRESSION_IF:
			return if_expression_sprint(&expr->cond, str);
		case EXPRESSION_FUNC:
			return func_literal_sprint(&expr->func, str);
		case EXPRESSION_CALL:
			return call_expression_sprint(&expr->call, str);
		case EXPRESSION_IDENT:
		case EXPRESSION_INT:
		case EXPRESSION_BOOL:
			return slice_sprint(&expr->token.literal, str);
	}
};

char *
statement_sprint(struct statement *st, char *str)
{
	switch (st->type) {
	case STATEMENT_LET:
		return let_statement_sprint(&st->let, str);
	case STATEMENT_RETURN:
		return return_statement_sprint(&st->retrn, str);
	case STATEMENT_EXPRESSION:
		return expression_statement_sprint(&st->expr, str);
	case STATEMENT_BLOCK:
		return block_statement_sprint(&st->block, str);
	}
}

char *
program_sprint(struct program *prog, char *str)
{
	size_t i;
	struct statement *st;
	char *cur = str;
	vector_foreach(prog->statements, i, st) {
		statement_sprint(st, cur);
		cur += strlen(cur);
	}
	return str;
}

struct vector *
expression_vector_dup(const struct vector *vec)
{
	struct vector *dupv = vector_new_with_cap(vec->cap);
	size_t i;
	struct expression *expr;
	vector_foreach(vec, i, expr) {
		struct expression *dupexpr = node_dup(expr);
		vector_push(dupv, dupexpr);
	}

	return dupv;
}

struct vector *
statement_vector_dup(const struct vector *vec)
{
	struct vector *dupv = vector_new_with_cap(vec->cap);
	size_t i;
	struct statement *subst;
	vector_foreach(vec, i, subst) {
		struct statement *dupsub = node_dup(subst);
		vector_push(dupv, dupsub);
	}

	return dupv;
}

struct expression *
expression_dup(const struct expression *expr)
{
	MEMDUP(struct expression *, dupd, expr);
	switch (expr->type) {
	case EXPRESSION_IDENT:
	case EXPRESSION_INT:
	case EXPRESSION_BOOL:
		return dupd;
	case EXPRESSION_PREFIX:
		dupd->prefix.right = node_dup(expr->prefix.right);
		return dupd;
	case EXPRESSION_INFIX:
		dupd->infix.left = node_dup(expr->infix.left);
		dupd->infix.right = node_dup(expr->infix.right);
		return dupd;
	case EXPRESSION_IF:
		dupd->cond.condition = node_dup(expr->cond.condition);
		dupd->cond.consequence = node_dup(expr->cond.consequence);
		dupd->cond.alternative = node_dup(expr->cond.alternative);
		return dupd;
	case EXPRESSION_FUNC:
		dupd->func.parameters = expression_vector_dup(expr->func.parameters);
		dupd->func.body = node_dup(expr->func.body);
		return dupd;
	case EXPRESSION_CALL:
		dupd->call.func = node_dup(expr->call.func);
		dupd->call.arguments = expression_vector_dup(expr->call.arguments);
		return dupd;
	}
}

struct statement *
statement_dup(const struct statement *st)
{
	MEMDUP(struct statement *, dupd, st);
	switch (st->type) {
	case STATEMENT_LET:
		MEMDUP(, dupd->let.name, st->let.name);
		dupd->let.value = node_dup(st->let.value);
		return dupd;
	case STATEMENT_RETURN:
		dupd->retrn.value = node_dup(st->retrn.value);
		return dupd;
	case STATEMENT_EXPRESSION:
		dupd->expr.expr = node_dup(st->expr.expr);
		return dupd;
	case STATEMENT_BLOCK:
		dupd->block.statements = statement_vector_dup(st->block.statements);
		return dupd;
	}
}

void
expression_destroy(struct expression *expr)
{
	switch (expr->type) {
	case EXPRESSION_PREFIX:
		expression_destroy(expr->prefix.right);
		break;
	case EXPRESSION_INFIX:
		expression_destroy(expr->infix.right);
		expression_destroy(expr->infix.left);
		break;
	case EXPRESSION_IF:
		expression_destroy(expr->cond.condition);
		statement_destroy(expr->cond.consequence);
		if (expr->cond.alternative) statement_destroy(expr->cond.alternative);
		break;
	case EXPRESSION_FUNC: {
		size_t i;
		struct expression *param;
		vector_foreach(expr->func.parameters, i ,param) {
			node_destroy(param);
		}
		vector_free(expr->func.parameters);
		node_destroy(expr->func.body);
		break;
	}
	case EXPRESSION_CALL: {
		size_t i;
		struct expression *arg;
		vector_foreach(expr->call.arguments, i ,arg) {
			node_destroy(arg);
		}
		expression_destroy(expr->call.func);
		vector_free(expr->call.arguments);
		break;
	}
	default:
		break;
	}
	free(expr);
}

void
statement_destroy(struct statement *st)
{
	switch (st->type) {
	case STATEMENT_LET:
		free(st->let.name);
		node_destroy(st->let.value);
		break;
	case STATEMENT_RETURN:
		node_destroy(st->retrn.value);
		break;
	case STATEMENT_EXPRESSION:
		if (st->expr.expr) node_destroy(st->expr.expr);
		break;
	case STATEMENT_BLOCK:;
		size_t i;
		struct statement *subst;
		vector_foreach(st->block.statements, i ,subst) {
			node_destroy(subst);
		}
		vector_free(st->block.statements);
		break;
	}
	free(st);
}

void
program_destroy(struct program *prog)
{
	size_t i;
	struct statement *st;
	vector_foreach(prog->statements, i ,st) {
		node_destroy(st);
	}
	vector_free(prog->statements);
	free(prog);
}