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

#define VEC_NAME tokens
#define VEC_TYPE char *
#include <conts/vec.h>

struct def {
	char *id;
	char *types;
	char *args;
	char *inst;
	char *body;
};

#define MAP_NAME defs
#define MAP_KEY char *
#define MAP_TYPE struct def
#define MAP_CMP(a, b) strcmp((a), (b))
#define MAP_HASH(a) CONTS_MAP_STR_HASH(a)
#include <conts/map.h>

struct defs templates;
struct defs traits;

static void usage(FILE *f)
{
	fprintf(f, "usage: ngc <file>\n");
}

static char *append(char *grp, char c)
{
	/* pretty slow but good enough for now I guess */
	size_t len = strlen(grp);
	grp = realloc(grp, len + 2);
	grp[len] = c;
	grp[len + 1] = '\0';
	return grp;
}

static bool char_in(char c, const char *group)
{
	for (size_t i = 0; i < strlen(group); ++i)
		if (c == group[i])
			return true;

	return false;
}

static void optional_str_space(char *str, size_t *idx)
{
	if (!isspace(str[*idx]))
		return;

	while (str[*idx] && isspace(str[*idx])) {
		(*idx)++;
	}
}

static char *next_str_grp(char *grp, char *str, size_t *idx)
{
	grp = realloc(grp, sizeof(char));
	grp[0] = '\0';

	while (1) {
		char c = str[*idx];
		if (!c)
			return grp;

		if (isspace(c))
			return grp;

		/* special separators */
		if (char_in(c, "{[()]},.;"))
			return grp;

		grp = append(grp, c);
		(*idx)++;
	}

	free(grp);
	return NULL;
}

static char *next_str_until(char *grp, char *str, size_t *idx, char *until)
{
	grp = realloc(grp, sizeof(char));
	grp[0] = '\0';

	while (1) {
		char c = str[*idx];
		if (!c)
			return grp;

		if (char_in(c, until))
			return grp;

		grp = append(grp, c);
		(*idx)++;
	}

	free(grp);
	return NULL;
}

static char *optional_space(FILE *f)
{
	int c = fgetc(f);
	if (c == EOF)
		return NULL;

	if (!isspace(c)) {
		ungetc(c, f);
		return NULL;
	}

	char *space = malloc(2 * sizeof(char));
	space[0] = c;
	space[1] = '\0';

	while (1) {
		c = fgetc(f);
		if (c == EOF)
			return space;

		if (!isspace(c)) {
			ungetc(c, f);
			return space;
		}

		space = append(space, c);
	}

	return NULL;
}

static char *next_grp(char *grp, FILE *f)
{
	/* the worst C parser you'll ever see, lol */
	int c = fgetc(f);
	if (c == EOF) {
		free(grp);
		return NULL;
	}

	grp = realloc(grp, 2 * sizeof(char *));
	grp[0] = c;
	grp[1] = '\0';

	if (isspace(c)) {
		while (1) {
			c = fgetc(f);
			if (c == EOF)
				return grp;

			if (!isspace(c)) {
				ungetc(c, f);
				return grp;
			}

			grp = append(grp, c);
		}

		return NULL;
	}

	while (1) {
		c = fgetc(f);
		if (c == EOF)
			return grp;

		if (isspace(c)) {
			ungetc(c, f);
			return grp;
		}

		/* special separators */
		if (char_in(c, "{[(,.*")) {
			ungetc(c, f);
			return grp;
		}

		grp = append(grp, c);
		if (char_in(c, "}]);"))
			return grp;
	}

	return NULL;
}

static void dump_tokens(struct tokens *tokens)
{
	foreach(tokens, t, tokens) {
		printf("%s", *t);
		free(*t);
	}
	tokens_destroy(tokens);
}

static void drop_tokens(struct tokens *tokens)
{
	foreach(tokens, t, tokens) {
		free(*t);
	}
	tokens_destroy(tokens);
}

static char *collect_until(char *grp, FILE *f, int until)
{
	grp = realloc(grp, sizeof(char));
	grp[0] = '\0';

	while (1) {
		int c = fgetc(f);
		assert(c != EOF);

		if (c == until) {
			ungetc(c, f);
			return grp;
		}

		grp = append(grp, c);
	}

	return NULL;
}

static char *collect_block(char *grp, FILE *f, int start, int end)
{
	int c = fgetc(f);
	assert(c == start);

	grp = realloc(grp, 2 * sizeof(char));
	grp[0] = c;
	grp[1] = '\0';

	size_t depth = 1;
	while (1) {
		int c = fgetc(f);
		assert(c != EOF);
		grp = append(grp, c);

		if (c == start)
			depth++;

		if (c == end)
			depth--;

		if (depth == 0)
			return grp;
	}

	return NULL;
}

static char *collect_str_block(char *grp, char *str, size_t *idx, int start, int end)
{
	char c = str[(*idx)++];
	assert(c == start);

	grp = realloc(grp, 2 * sizeof(char));
	grp[0] = start;
	grp[1] = '\0';

	size_t depth = 1;
	while (1) {
		c = str[(*idx)++];
		grp = append(grp, c);

		if (c == start)
			depth++;

		if (c == end)
			depth--;

		if (depth == 0)
			return grp;

	}

	return NULL;
}

static int is_identifier(char *id)
{
	while (*id) {
		if (!isalnum(*id) && *id != '_')
			return false;

		id++;
	}

	return true;
}

static char *collect_square_block(char *grp, FILE *f)
{
	return collect_block(grp, f, '[', ']');
}

static char *collect_paren_block(char *grp, FILE *f)
{
	return collect_block(grp, f, '(', ')');
}

static char *collect_curly_block(char *grp, FILE *f)
{
	return collect_block(grp, f, '{', '}');
}

static char *collect_str_square_block(char *grp, char *str, size_t *idx)
{
	return collect_str_block(grp, str, idx, '[', ']');
}

static char *collect_str_paren_block(char *grp, char *str, size_t *idx)
{
	return collect_str_block(grp, str, idx, '(', ')');
}

static char *process_trait(char *id, struct tokens *tokens, char *grp, FILE *f)
{
	printf("/* found trait %s*/", id);
	grp = collect_curly_block(grp, f);
	assert(grp[0] == '{');

	struct def *exists = defs_find(&traits, id);
	if (exists) {
		fprintf(stderr, "trait %s redefined\n", id);
		drop_tokens(tokens);
		return grp;
	}

	struct def def = {
		.id = strdup(id),
		.types = NULL,
		.args = NULL,
		.body = strdup(grp)
	};

	defs_insert(&traits, def.id, def);
	drop_tokens(tokens);
	return grp;
}

static char *process_supertemplate(char *id, char *types, char *args, struct tokens *tokens, char *grp, FILE *f)
{
	printf("/* found supertemplate %s */", id);

	/* discard equals sign */
	int c = fgetc(f);
	assert(c == '=');

	grp = collect_until(grp, f, '{');
	assert(grp);
	char *inst = strdup(grp);
	tokens_append(tokens, inst);

	grp = collect_curly_block(grp, f);
	char *body = strdup(grp);
	tokens_append(tokens, body);

	char *space = NULL;
	if ((space = optional_space(f)))
		tokens_append(tokens, space);

	struct def *exists = defs_find(&templates, id);
	if (exists) {
		fprintf(stderr, "template %s redefined\n", id);
		drop_tokens(tokens);
		return grp;
	}

	struct def def = {
		.id = strdup(id),
		.types = strdup(types),
		.args = strdup(args),
		.body = strdup(body),
		.inst = strdup(inst)
	};

	defs_insert(&templates, def.id, def);
	drop_tokens(tokens);
	return grp;
}

static char *maybe_process_template(char *id, struct tokens *tokens, char *grp, FILE *f)
{
	printf("/* maybe found template */");
	/* [type args] */
	grp = collect_square_block(grp, f);
	char *types = strdup(grp);
	tokens_append(tokens, types);

	char *space = optional_space(f);
	if (space)
		tokens_append(tokens, space);

	/* should be () if we're a template */
	int c = fgetc(f);
	ungetc(c, f);
	if (c != '(') {
		printf("/* nope, no parens */");
		dump_tokens(tokens);
		return grp;
	}

	/* (value args) */
	grp = collect_paren_block(grp, f);
	char *args = strdup(grp);
	tokens_append(tokens, args);

	space = optional_space(f);
	if (space)
		tokens_append(tokens, space);

	c = fgetc(f);
	ungetc(c, f);
	if (c == '=')
		return process_supertemplate(id, types, args, tokens, grp, f);

	if (c != '{') {
		printf("/* nope, no block */");
		dump_tokens(tokens);
		return grp;
	}

	printf("/* it is! %s*/", id);
	char *body = grp = collect_curly_block(grp, f);

	struct def *exists = defs_find(&templates, id);
	if (exists) {
		fprintf(stderr, "template %s redefined\n", id);
		drop_tokens(tokens);
		return grp;
	}

	struct def def = {
		/* these could be gotten rid of but works for now */
		.id = strdup(id),
		.types = strdup(types),
		.args = strdup(args),
		.body = strdup(body),
		.inst = NULL
	};

	defs_insert(&templates, def.id, def);
	drop_tokens(tokens);
	return grp;
}

static char *replace_id(char *body, char *replace, char *id)
{
	assert(body);
	assert(id);
	size_t len   = strlen(replace);
	size_t idlen = strlen(id);

	/* check that replacement is regular alphanumeric */
	for (size_t i = 0; i < len; ++i) {
		assert(isalnum(replace[i]) || replace[i] == '_');
	}

	char *new = malloc(1 * sizeof(char));
	new[0] = '\0';

	size_t idx = 0;
	while (body[idx]) {
		if (!isalnum(body[idx]) && body[idx] != '_') {
			new = append(new, body[idx++]);
			continue;
		}

		size_t start = idx;
		while (body[idx] && (isalnum(body[idx]) || body[idx] == '_'))
			idx++;

		size_t diff = idx - start;
		if (strncmp(&body[start], replace, diff) == 0) {
			for (size_t i = 0; i < idlen; ++i)
				new = append(new, id[i]);

			continue;
		}

		for (size_t i = 0; i < diff; ++i)
			new = append(new, body[start + i]);
	}

	free(body);
	return new;
}

static char *replace(char *body, char *replace, char *id)
{
	assert(body);
	assert(id);

	char *new = malloc(1 * sizeof(char));
	new[0] = '\0';

	size_t len   = strlen(replace);
	size_t idlen = strlen(id);

	size_t idx = 0;
	while (body[idx]) {
		if (body[idx] != replace[0]) {
			new = append(new, body[idx++]);
			continue;
		}

		if (strncmp(&body[idx], replace, len) == 0) {
			for (size_t i = 0; i < idlen; ++i)
				new = append(new, id[i]);

			idx += len;
			continue;
		}
	}

	free(body);
	return new;
}

static void expand_instance(char *id, char *instance)
{
	char *grp = NULL;

	size_t instance_idx = 0;
	optional_str_space(instance, &instance_idx);
	grp = next_str_grp(grp, instance, &instance_idx);
	assert(is_identifier(grp));
	char *name = strdup(grp);

	optional_str_space(instance, &instance_idx);
	grp = collect_str_square_block(grp, instance, &instance_idx);
	char *types = strdup(grp);

	optional_str_space(instance, &instance_idx);
	grp = collect_str_paren_block(grp, instance, &instance_idx);
	char *args = strdup(grp);

	optional_str_space(instance, &instance_idx);
	struct def *def = defs_find(&templates, name);
	if (!def) {
		/** @todo location info and so on */
		fprintf(stderr, "no such template: %s\n", name);
		free(name);
		free(args);
		free(types);
		return;
	}

	free(name);


	/** @todo check traits */

	/* +1 to skip leading '{' */
	char *body = strdup(def->body + 1);
	/* overwrite trailing '}' */
	body[strlen(body) - 1] = '\0';

	/* handle type args */
	char *params = def->types;
	assert(params[0] == '[');

	/* skip '[' */
	size_t params_idx = 1;

	/* skip '(' in args */
	size_t types_idx = 1;

	struct tokens type_params = tokens_create(0);
	struct tokens type_args   = tokens_create(0);

	while (1) {
		/* get type param */
		optional_str_space(params, &params_idx);

		grp = next_str_grp(grp, params, &params_idx);
		assert(is_identifier(grp));
		char *trait = strdup(grp);

		optional_str_space(params, &params_idx);
		grp = next_str_grp(grp, params, &params_idx);
		assert(is_identifier(grp));
		char *name = strdup(grp);

		optional_str_space(params, &params_idx);
		char delim = params[params_idx++];
		assert(delim == ',' || delim == ']');

		/* get arg type */
		optional_str_space(types, &types_idx);
		grp = next_str_until(grp, types, &types_idx, ",]");
		char *type = strdup(grp);

		char last = types[types_idx++];
		assert(last == ',' || last == ']');

		tokens_append(&type_params, name);
		tokens_append(&type_args,   type);

		free(trait);

		if (delim == ']' && last == ']')
			break;

		if (delim == ']') {
			fprintf(stderr, "more type args than params\n");
			free(body);
			drop_tokens(&type_params);
			drop_tokens(&type_args);
			free(args);
			free(types);
			free(grp);
			return;
		}

		if (last == ']') {
			fprintf(stderr, "more type params than args\n");
			free(body);
			drop_tokens(&type_params);
			drop_tokens(&type_args);
			free(args);
			free(types);
			free(grp);
			return;
		}
	}

	char *super = NULL;
	body        = replace(body, "<>", id);

	if (def->inst)
		super = replace(strdup(def->inst), "<>", id);

	for (size_t i = 0; i < tokens_len(&type_params); ++i) {
		char *name = *tokens_at(&type_params, i);
		char *type = *tokens_at(&type_args,   i);

		body  = replace_id(body, name, type);

		if (super)
			super = replace_id(super, name, type);
	}

	if (super)
		expand_instance(id, super);

	/** @todo regular params, probably much like above but parsing the type
	 * might be a bit annoying... */

	drop_tokens(&type_params);
	drop_tokens(&type_args);

	printf(body);
	free(body);
	free(super);

	free(types);
	free(args);
	free(grp);
}

static char *process_instance(char *id, struct tokens *tokens, char *grp, FILE *f)
{
	printf("/* found instance %s */", id);
	/* discard equals sign */
	int c = fgetc(f);
	assert(c == '=');

	grp = collect_until(grp, f, ';');
	char *instance = strdup(grp);
	tokens_append(tokens, instance);

	expand_instance(id, instance);

	/* ; can probably be ignored...?*/
	drop_tokens(tokens);
	return grp;
}

static char *process_typedef(char *grp, FILE *f)
{
	struct tokens tokens = tokens_create(3);
	tokens_append(&tokens, strdup(grp));

	/* space */
	grp = next_grp(grp, f);
	assert(grp);
	tokens_append(&tokens, strdup(grp));

	/* id */
	grp = next_grp(grp, f);
	assert(grp);

	char *id = strdup(grp);
	assert(is_identifier(id));
	tokens_append(&tokens, id);

	/* space/check */
	char *space = optional_space(f);
	if (space)
		tokens_append(&tokens, space);

	char c = fgetc(f);
	ungetc(c, f);
	if (c == '{'
		&& strcmp(id, "struct") != 0
		&& strcmp(id, "union") != 0
		&& strcmp(id, "enum") != 0)
		return process_trait(id, &tokens, grp, f);

	if (c == '[')
		return maybe_process_template(id, &tokens, grp, f);

	if (c == '=')
		return process_instance(id, &tokens, grp, f);

	dump_tokens(&tokens);
	return grp;
}

static int process(FILE *f)
{
	char *grp = NULL;
	while ((grp = next_grp(grp, f))) {
		if (strcmp(grp, "typedef") == 0)
			grp = process_typedef(grp, f);
		else
			printf("%s", grp);
	}

	return 0;
}

static void free_defs(struct defs *defs)
{
	foreach(defs, def, defs) {
		struct def d = def->data;
		free(d.id);
		free(d.types);
		free(d.args);
		free(d.body);
		free(d.inst);
	}

	defs_destroy(defs);
}

int main(int argc, char *argv[])
{
	if (argc != 2) {
		usage(stderr);
		return -1;
	}

	templates = defs_create(0);
	traits = defs_create(0);

	char *fname = argv[1];

	FILE *f = NULL;
	if (strcmp(fname, "-") == 0)
		f = stdin;
	else
		f = fopen(fname, "rb");

	assert(f);

	int ret = process(f);

	free_defs(&templates);
	free_defs(&traits);
	fclose(f);
	return ret;
}