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

#include <tree_sitter/api.h>

#include "gen/gen_enum_tsq.h"

TSLanguage const *tree_sitter_c(void);

char const *USAGE = "\
cgen <mode>\
";

size_t read_stdin(char *buf, size_t buf_cap) {
	char *at = buf;
	char *end = buf + buf_cap;

	for (; at < end; ++at) {
		if ((*at = fgetc(stdin)) == EOF) {
			break;
		}
	}
	if (at < end) { *at = '\0'; }

	return at - buf;
}

void gen_enum_strs() {
	size_t buf_cap = 2 << 16;
	char *buf = calloc(buf_cap, sizeof(*buf));
	size_t buf_len = read_stdin(buf, buf_cap);

	TSParser *parser = ts_parser_new();
	TSLanguage const *lang_c = tree_sitter_c();
	ts_parser_set_language(parser, lang_c);

	TSTree *tree = ts_parser_parse_string(parser, nullptr, buf, buf_len);
	TSNode root_node = ts_tree_root_node(tree);

	uint32_t err_offset = 0;
	TSQueryError err = {};
	TSQuery *query = ts_query_new(
		lang_c,
		(char*)src_gen_enum_tsq,
		src_gen_enum_tsq_len,
		&err_offset,
		&err
	);
	if (err != TSQueryErrorNone) {
		printf("query error %d near: %s\n", err, src_gen_enum_tsq + err_offset);
		assert(false);
	}

	uint32_t cap_id_poison = 2 << 16;
	uint32_t cap_id_enum_name = cap_id_poison;
	uint32_t cap_id_enum_variant = cap_id_poison;
	uint32_t cap_id_enum_declarator_name = cap_id_poison;
	uint32_t cap_count = ts_query_capture_count(query);
	for (uint32_t i = 0; i < cap_count; ++i) {
		uint32_t len = 0;
		char const *name = ts_query_capture_name_for_id(query, i, &len);

		if (!strncmp(name, "enum_name", len)) {
			cap_id_enum_name = i;
		}

		if (!strncmp(name, "enum_variant", len)) {
			cap_id_enum_variant = i;
		}

		if (!strncmp(name, "enum_declarator_name", len)) {
			cap_id_enum_declarator_name = i;
		}
	}

	TSQueryCursor *qc = ts_query_cursor_new();
	ts_query_cursor_exec(qc, query, root_node);

	TSQueryMatch match = {};
	while (ts_query_cursor_next_match(qc, &match)) {
		char const *poison = "POISON";
		char const *enum_name = poison;
		size_t enum_name_len = strlen(poison);

		size_t enum_variants_cap = 512;
		struct {
			char const *start;
			size_t len;
		} enum_variants[enum_variants_cap] = {};
		size_t enum_variants_len = 0;
		for (uint32_t i = 0; i < match.capture_count; ++i) {
			TSQueryCapture cap = match.captures[i];

			uint32_t start = ts_node_start_byte(cap.node);
			uint32_t end = ts_node_end_byte(cap.node);
			uint32_t len = end - start;
			char const *at = buf + start;

			if (
				cap.index == cap_id_enum_name
				|| cap.index == cap_id_enum_declarator_name
			) {
				enum_name = at;
				enum_name_len = len;
			}

			if (
				cap.index == cap_id_enum_variant
				&& enum_variants_len < enum_variants_cap
			) {
				enum_variants[enum_variants_len].len = len;
				enum_variants[enum_variants_len].start = at;

				enum_variants_len++;
			}
		}

		printf("char const *%.*s_strs[] = {\n", (int)enum_name_len, enum_name);
		for (size_t i = 0; i < enum_variants_len; ++i) {
			typeof(enum_variants[0]) evar = enum_variants[i];
			printf("\t[%.*s] = \"%.*s\",\n", (int)evar.len, evar.start, (int)evar.len, evar.start);
		}
		printf("};\n");
	}
}

int main(int argc, char **argv) {
	if (argc != 2) {
		printf("%s", USAGE);
		exit(1);
	}

	char const *mode = argv[1];
	if (!strcmp(mode, "enum-strs")) {
		gen_enum_strs();
	}

	return 0;
}