#include <assert.h>
#include <stdio.h>
#include <stdlib.h>

#include <log.h>
#include <tokenize.h>
#include <treeify.h>

#define INITIAL_PARAMETER_COUNT 3

static void __expect(struct ctx *ctx, struct token *token, enum token_type type, const char *file,
		     int line, const char *func, const char *type_enum)
{
	if (token->type != type) {
		printf("[DBG] %s:%d: %s\n", file, line, func);
		token_print(ctx, token);
		errln(&token->location, "Expected token of type %s", type_enum);
	}
}

#define expect(token, type) __expect(ctx, token, type, __FILE__, __LINE__, __func__, #type)

static struct token *next(struct token *token, size_t i)
{
	return token + i;
}

static struct token *parse_declaration(struct ctx *ctx, struct token *token)
{
	expect(token, IDENT);

	struct node_declaration *node = malloc(sizeof(*node));
	node->callee.name = token->string;

	node->parameters = malloc(INITIAL_PARAMETER_COUNT * sizeof(*node->parameters));
	size_t param_idx = 0;

	token = next(token, 1);
	while (token->type != TYPEDELIM) {
		if (token->type == NEWLINE || token->type == END)
			expect(token, TYPEDELIM);

		if (token->type != TYPE)
			expect(token, TYPE);

		// Expand parameter space if necessary
		if ((param_idx + 1) % INITIAL_PARAMETER_COUNT == 0)
			// TODO: Fix realloc failure check (and other mallocs too btw)
			node->parameters = realloc(node->parameters,
						   ((param_idx / INITIAL_PARAMETER_COUNT) + 1) *
							   INITIAL_PARAMETER_COUNT *
							   sizeof(*node->parameters));

		node->parameters[param_idx].type = token->string;
		param_idx++;

		token = next(token, 1);
	}

	node->parameter_count = param_idx;

	token = next(token, 1);
	expect(token, TYPE);
	node->callee.type = token->string;

	tree_add(ctx, DECLARATION, node); // TODO: Push to declaration/signature array instead

	expect(next(token, 1), NEWLINE);
	return next(token, 2);
}

static struct token *parse_definition(struct ctx *ctx, struct token *token)
{
	expect(token, IDENT);

	struct node_definition *node = malloc(sizeof(*node));
	node->callee.name = token->string;

	node->parameters = malloc(INITIAL_PARAMETER_COUNT * sizeof(*node->parameters));
	size_t param_idx = 0;

	token = next(token, 1);
	while (token->type != IDENTDELIM) {
		if (token->type == NEWLINE || token->type == END)
			expect(token, IDENTDELIM);

		if (token->type != PARAM)
			expect(token, PARAM);

		// Expand parameter space if necessary
		if ((param_idx + 1) % INITIAL_PARAMETER_COUNT == 0)
			// TODO: Fix realloc failure check (and other mallocs too btw)
			node->parameters = realloc(node->parameters,
						   ((param_idx / INITIAL_PARAMETER_COUNT) + 1) *
							   INITIAL_PARAMETER_COUNT *
							   sizeof(*node->parameters));

		node->parameters[param_idx].name = token->string;
		param_idx++;

		token = next(token, 1);
	}

	node->parameter_count = param_idx;

	tree_add(ctx, DEFINITION, node);

	// TODO: Parse expression
	while (token->type != NEWLINE)
		token = next(token, 1);

	return next(token, 1);
}

static struct token *parse_block(struct ctx *ctx, struct token *token)
{
	if (token->type != IDENT)
		return next(token, 1); //&ctx->tokens[ctx->token_count - 1];

	token = parse_declaration(ctx, token);
	token = parse_definition(ctx, token);
	expect(token, NEWLINE);
	return next(token, 1);
}

static struct token *parse(struct ctx *ctx, struct token *token)
{
	while (token->type != END)
		token = parse_block(ctx, token);

	return token;
}

struct node *tree_create(void)
{
	struct node *tree = malloc(sizeof(*tree));
	tree->prev = NULL;
	tree->next = NULL;
	return tree;
}

void tree_add(struct ctx *ctx, enum node_type type, void *data)
{
	assert(ctx->tree.head);
	struct node *node = calloc(sizeof(*node), 1);
	node->type = type;
	node->data = data;
	if (!ctx->tree.current) {
		ctx->tree.head->next = node;
		node->prev = ctx->tree.head;
	} else {
		ctx->tree.current->next = node;
		node->prev = ctx->tree.current;
	}
	ctx->tree.current = node;
}

void tree_destroy(struct node *node)
{
	while (node) {
		struct node *next = node->next;

		if (node->type == DEFINITION) {
			struct node_definition *definition = node->data;
			free(definition->parameters);
			free(definition);
		} else if (node->type == DECLARATION) {
			struct node_declaration *declaration = node->data;
			free(declaration->parameters);
			free(declaration);
		}

		free(node);
		node = next;
	}
}

void treeify(struct ctx *ctx)
{
	struct token *token = ctx->tokens;

	while (token->type == NEWLINE)
		token = next(token, 1);

	token = parse(ctx, token);
	expect(token, END);
}