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

#include "nbt.h"
#include "lib.h"

typedef enum {
	TOK_RBRACE,
	TOK_LBRACE,
	TOK_RBRACK,
	TOK_LBRACK,
	TOK_COMMA,
	TOK_COLON,
	TOK_SEMI_COLON,
	TOK_STRING,
	TOK_IDENT,
} tokentype_t;

typedef union {
	struct {
		char *data;
		uint16_t len;
	} text;
	double d;
	int64_t l;
} tokendata_t;

typedef struct {
	tokentype_t type;
	tokendata_t data;
} token_t;

static char ret = '\0';

static void snbt_free_token(token_t *token) {
	if (token->type == TOK_IDENT || token->type == TOK_STRING) {
		free(token->data.text.data);
	}
}

static bool snbt_parse_string(token_t *token, const stream_t *stream, char delimiter) {
	int len = 0;
	int capacity = 8;
	char *buf = xalloc(capacity * sizeof(char));

	while (1) {
		char c;
		if (stream_read(stream, &c, 1) == false) {
			free(buf);
			return false;
		}

		if (c == delimiter)
			break;

		if (c == '\\') {
			if (stream_read(stream, &c, 1) == false) {
				free(buf);
				return false;
			}

			if (c != '\\' && c != delimiter) {
				free(buf);
				return false;
			}
		}

		if (len == capacity) {
			capacity *= 2;
			buf = xrealloc(buf, capacity);
		}

		buf[len++] = c;
	}

	token->data.text.data = xalloc(len * sizeof(char));
	token->data.text.len = len;
	memcpy(token->data.text.data, buf, len * sizeof(char));
	free(buf);

	return true;
}

bool snbt_allowed_ident(char c) {
	if (c >= '0' && c <= '9')
		return true;
	if (c >= 'a' && c <= 'z')
		return true;
	if (c >= 'A' && c <= 'Z')
		return true;
	if (c == '_' || c == '-' || c == '.' || c == '+')
		return true;
	return false;
}

static bool snbt_parse_ident(token_t *token, const stream_t *stream, char first) {
	int len = 0;
	int capacity = 8;
	char *buf = xalloc(capacity * sizeof(char));

	buf[len++] = first;

	while (1) {
		char c;
		if (stream_read(stream, &c, 1) == false) {
			free(buf);
			return false;
		}

		if (snbt_allowed_ident(c) == false) {
			if (len == 0) {
				free(buf);
				return false;
			} else {
				ret = c;
				break;
			}
		}

		if (len == capacity) {
			capacity *= 2;
			buf = xrealloc(buf, capacity);
		}

		buf[len++] = c;
	}

	token->data.text.data = xalloc(len * sizeof(char));
	token->data.text.len = len;
	memcpy(token->data.text.data, buf, len * sizeof(char));
	free(buf);

	return true;
}

static bool snbt_next_token(token_t *token, const stream_t *stream) {

	memset(token, 0, sizeof(token_t));

	char c;
	bool ok = true;

retry:

	c = ret;
	ret = '\0';

	if (c == '\0' && stream_read(stream, &c, 1) == false)
		return false;

	switch (c) {
		case ' ':
		case '\t':
		case '\n':
		case '\r':
			goto retry;
		case '{':
			token->type = TOK_LBRACE;
			break;
		case '}':
			token->type = TOK_RBRACE;
			break;
		case '[':
			token->type = TOK_LBRACK;
			break;
		case ']':
			token->type = TOK_RBRACK;
			break;
		case ':':
			token->type = TOK_COLON;
			break;
		case ';':
			token->type = TOK_SEMI_COLON;
			break;
		case ',':
			token->type = TOK_COMMA;
			break;
		case '"':
		case '\'':
			token->type = TOK_STRING;
			ok = snbt_parse_string(token, stream, c);
			break;
		default:
			token->type = TOK_IDENT;
			ok = snbt_parse_ident(token, stream, c);
			break;
	}

	return ok;
}

static bool snbt_read_value(tag_t *tag, const stream_t *stream, token_t *first);
static bool snbt_convert_ident(tag_t *tag, token_t token);

static bool snbt_read_byte_array(tag_t *tag, const stream_t *stream) {
	int capacity = 8;
	int len = 0;
	int8_t *buf = xalloc(capacity * sizeof(int8_t));

	while(1) {
		tag_t value;
		if (snbt_read_value(&value, stream, NULL) == false) {
			free(buf);
			return false;
		}

		if (value.type != TAG_BYTE) {
			tag_free(&value);
			free(buf);
			return false;
		}
		if (len == capacity) {
			capacity *= 2;
			buf = xrealloc(buf, capacity * sizeof(int8_t));
		}
		buf[len++] = value.data.b;
		token_t token;
		if (snbt_next_token(&token, stream) == false) {
			free(buf);
			return false;
		}
		if (token.type == TOK_COMMA)
			continue;
		if (token.type == TOK_RBRACK)
			break;
		free(buf);
		return false;
	}
	tag->type = TAG_BYTE_ARRAY;
	tag->data.b_arr.data = xalloc(len * sizeof(int8_t));
	tag->data.b_arr.size = len;
	memcpy(tag->data.b_arr.data, buf, len * sizeof(int8_t));
	free(buf);

	return true;
}

static bool snbt_read_int_array(tag_t *tag, const stream_t *stream) {
	int capacity = 8;
	int len = 0;
	int32_t *buf = xalloc(capacity * sizeof(int32_t));

	while(1) {
		tag_t value;
		if (snbt_read_value(&value, stream, NULL) == false) {
			free(buf);
			return false;
		}
		if (value.type != TAG_INT) {
			tag_free(&value);
			free(buf);
			return false;
		}
		if (len == capacity) {
			capacity *= 2;
			buf = xrealloc(buf, capacity * sizeof(int32_t));
		}
		buf[len++] = value.data.i;
		token_t token;
		if (snbt_next_token(&token, stream) == false) {
			free(buf);
			return false;
		}
		if (token.type == TOK_COMMA)
			continue;
		if (token.type == TOK_RBRACK)
			break;
		free(buf);
		return false;
	}
	tag->type = TAG_INT_ARRAY;
	tag->data.i_arr.data = xalloc(len * sizeof(int32_t));
	tag->data.i_arr.size = len;
	memcpy(tag->data.i_arr.data, buf, len * sizeof(int32_t));
	free(buf);

	return true;
}

static bool snbt_read_long_array(tag_t *tag, const stream_t *stream) {
	int capacity = 8;
	int len = 0;
	int64_t *buf = xalloc(capacity * sizeof(int64_t));

	while(1) {
		tag_t value;
		if (snbt_read_value(&value, stream, NULL) == false) {
			free(buf);
			return false;
		}
		if (value.type != TAG_LONG) {
			tag_free(&value);
			free(buf);
			return false;
		}
		if (len == capacity) {
			capacity *= 2;
			buf = xrealloc(buf, capacity * sizeof(int64_t));
		}
		buf[len++] = value.data.l;
		token_t token;
		if (snbt_next_token(&token, stream) == false) {
			free(buf);
			return false;
		}
		if (token.type == TOK_COMMA)
			continue;
		if (token.type == TOK_RBRACK)
			break;
		free(buf);
		return false;
	}
	tag->type = TAG_LONG_ARRAY;
	tag->data.l_arr.data = xalloc(len * sizeof(int64_t));
	tag->data.l_arr.size = len;
	memcpy(tag->data.l_arr.data, buf, len * sizeof(int64_t));
	free(buf);

	return true;
}

static bool snbt_read_array(tag_t *tag, tagtype_t type, const stream_t *stream) {
	switch(type) {
		case TAG_BYTE:
			return snbt_read_byte_array(tag, stream);
		case TAG_INT:
			return snbt_read_int_array(tag, stream);
		case TAG_LONG:
			return snbt_read_long_array(tag, stream);
		default:
			return false;
	}
}

static bool snbt_read_list(tag_t *tag, tag_t *first, token_t *tok_sav, const stream_t *stream) {
	int capacity = 8;
	int len = 0;
	tag_t *tags = xalloc(sizeof(tag_t) * capacity);
	tagtype_t type = TAG_END;

	if (first != NULL) {
		first->name_len = 0;
		first->name = "";
		tags[len++] = *first;
		type = first->type;
		goto endcheck;
	}

	if (tok_sav == NULL) {
		if (snbt_next_token(tok_sav, stream) == false) {
			free(tags);
			return false;
		}
		if (tok_sav->type == TOK_RBRACK) {
			goto end;
		}
	}

	while(1) {

		tag_t value;
		token_t check;

		if (snbt_read_value(&value, stream, tok_sav) == false) {
			snbt_free_token(tok_sav);
			free(tags);
			return false;
		}

		tok_sav = NULL;

		if (type != TAG_END && value.type != type) {
			free(tags);
			return false;
		}

		if (len == capacity) {
			capacity *= 2;
			tags = xrealloc(tags, capacity * sizeof(tag_t));
		}

		tags[len++] = value;

endcheck:

		if (snbt_next_token(&check, stream) == false) {
			free(tags);
			return false;
		}

		if (check.type == TOK_COMMA)
			continue;

		if (check.type == TOK_RBRACK)
			break;
			
		free(tags);
		return false;
	}
end:
	tag->type = TAG_LIST;
	tag->data.list.type = type;
	tag->data.list.size = len;
	tag->data.list.tags = xalloc(len * sizeof(tag_t));
	memcpy(tag->data.list.tags, tags, len *sizeof(tag_t));
	free(tags);

	return true;
}

static bool snbt_read_collection(tag_t *tag, const stream_t *stream) {

	token_t first;

	if (snbt_next_token(&first, stream) == false)
		return false;

	if (first.type == TOK_RBRACK) {
		tag->type = TAG_LIST;
		tag->data.list.type = TAG_END;
		tag->data.list.size = 0;
		tag->data.list.tags = NULL;
		return true;
	}

	if (first.type == TOK_IDENT && first.data.text.len == 1) {
		tagtype_t type;
		char c = first.data.text.data[0];
		switch (c) {
			case 'B':
				type = TAG_BYTE;
				break;
			case 'I':
				type = TAG_INT;
				break;
			case 'L':
				type = TAG_LONG;
				break;
			default: {
				tag_t new;
				if (snbt_convert_ident(&new, first) == false) {
					snbt_free_token(&first);
					return false;
				}
				return snbt_read_list(tag, &new, NULL, stream);
			}
		}

		token_t second;
		if (snbt_next_token(&second, stream) == false) {
			snbt_free_token(&first);
			return false;
		}

		if (second.type == TOK_COMMA) {
			tag_t new;
			new.type = TAG_STRING;
			new.data.string.data = first.data.text.data;
			new.data.string.size = first.data.text.len;
			return snbt_read_list(tag, &new, NULL, stream);
		}
		
		snbt_free_token(&first);

		if (second.type == TOK_SEMI_COLON) {
			return snbt_read_array(tag, type, stream);
		}

		snbt_free_token(&second);
		return false;

	}

	return snbt_read_list(tag, NULL, &first, stream);
}

static bool snbt_read_compound(tagdata_t *data, const stream_t *stream) {
	map_t map;
	map_init(&map);

	token_t next = {0};
	if (snbt_next_token(&next, stream) == false) {
		snbt_free_token(&next);
		return false;
	}

	if (next.type == TOK_RBRACE) {
		data->compound = map;
		return true;
	}

	while (1) {

		if (next.type != TOK_STRING && next.type != TOK_IDENT) {
			map_free(&map);
			snbt_free_token(&next);
			return false;
		}

		char *name = next.data.text.data;
		int name_len = next.data.text.len;

		if (name_len < 1) {
			map_free(&map);
			free(name);
			return false;
		}
		
		if (snbt_next_token(&next, stream) == false || next.type != TOK_COLON) {
			map_free(&map);
			free(name);
			return false;
		}
		
		tag_t value;
		if (snbt_read_value(&value, stream, NULL) == false) {
			map_free(&map);
			free(name);
			return false;
		}

		value.name = name;
		value.name_len = name_len;

		map_put(&map, &value);

		if (snbt_next_token(&next, stream) == false) {
			map_free(&map);
			snbt_free_token(&next);
			return false;
		}

		if (next.type == TOK_COMMA) {
			if (snbt_next_token(&next, stream) == false) {
				map_free(&map);
				return false;
			}
			continue;
		} else if (next.type == TOK_RBRACE) {
			break;
		} else {
			map_free(&map);
			snbt_free_token(&next);
			return false;
		}

	}

	data->compound = map;

	return true;
}

static bool snbt_convert_decimal(tag_t *tag, char *text) {
	uint16_t len = strlen(text);

	char *end = NULL;

	double d = strtod(text, &end);

	char c = *end;
	bool check1 = end == &text[len - 1] && (
		c == 'f' ||
		c == 'F' ||
		c == 'd' ||
		c == 'D'
	);
	bool check2 = end == &text[len] && c == '\0';
	if (!check1 && !check2) {
		return false;
	}
	
	if (tag->type == TAG_FLOAT) {
		tag->data.f = (float)d;
	} else {
		tag->data.d = d;
	}

	return true;
}

static bool snbt_convert_int(tag_t *tag, char *text) {
	uint16_t len = strlen(text);

	char *end = NULL;

	uint64_t i = strtol(text, &end, 10);

	char c = *end;
	bool check1 = end == &text[len - 1] && (
		c == 'b' ||
		c == 'B' ||
		c == 's' ||
		c == 'S' ||
		c == 'l' ||
		c == 'L'
	);
	bool check2 = end == &text[len] && c == '\0';
	if (!check1 && !check2) {
		return false;
	}

	if (tag->type == TAG_BYTE) {
		tag->data.b = (uint8_t)i;
	} else if (tag->type == TAG_SHORT) {
		tag->data.s = (uint16_t)i;
	} else if (tag->type == TAG_INT) {
		tag->data.i = (uint32_t)i;
	} else {
		tag->data.l = i;
	}

	return true;
}

static bool snbt_convert_ident(tag_t *tag, token_t token) {
	uint16_t len = token.data.text.len + 1;
	char *text = xalloc(len * sizeof(char));
	memcpy(text, token.data.text.data, len - 1);
	text[len - 1] = '\0';

	if (len == 4 && memcmp(text, "true", 4) == 0) {
		tag->type = TAG_BYTE;
		tag->data.b = 1;
	} else if (len == 5 && memcmp(text, "false", 5) == 0) {
		tag->type = TAG_BYTE;
		tag->data.b = 0;
	}

	// try to parse as a number lol he he ha ha
	char end = text[len - 2];

	bool ok = false;

	switch (end) {
		case 'f':
		case 'F':
			tag->type = TAG_FLOAT;
			ok = snbt_convert_decimal(tag, text);
			break;
		case 'd':
		case 'D':
			tag->type = TAG_DOUBLE;
			ok = snbt_convert_decimal(tag, text);
			break;
		case 'b':
		case 'B':
			tag->type = TAG_BYTE;
			ok = snbt_convert_int(tag, text);
			break;
		case 's':
		case 'S':
			tag->type = TAG_SHORT;
			ok = snbt_convert_int(tag, text);
			break;
		case 'l':
		case 'L':
			tag->type = TAG_LONG;
			ok = snbt_convert_int(tag, text);
			break;
		default:
			break;
	}

	if (ok) {
		free(text);
		snbt_free_token(&token);
		return true;
	}

	// unkown type try long >> double >> string (fallback)
	
	tag->type = TAG_INT;
	if (snbt_convert_int(tag, text) == true) {
		free(text);
		snbt_free_token(&token);
		return true;
	}

	tag->type = TAG_DOUBLE;
	if (snbt_convert_decimal(tag, text) == true) {
		free(text);
		snbt_free_token(&token);
		return true;
	}

	tag->type = TAG_STRING;
	tag->data.string.data = token.data.text.data;
	tag->data.string.size = token.data.text.len;
	free(text);
	return true;
}

static bool snbt_read_value(tag_t *tag, const stream_t *stream, token_t *first) {
	
	token_t token;

	if (first != NULL)
		token = *first;
	else if (snbt_next_token(&token, stream) == false)
		return false;

	tag->name = "";
	tag->name_len = 0;

	bool ok = true;

	switch (token.type) {
        case TOK_RBRACK:
        case TOK_RBRACE:
        case TOK_COLON:
		case TOK_SEMI_COLON:
        case TOK_COMMA:
			ok = false;
			break;
        case TOK_LBRACK:
			tag->type = TAG_LIST;
			ok = snbt_read_collection(tag, stream);
			break;
        case TOK_LBRACE:
			tag->type = TAG_COMPOUND;
			ok = snbt_read_compound(&tag->data, stream);
			break;
        case TOK_STRING:
			tag->type = TAG_STRING;
			tag->data.string.data = token.data.text.data;
			tag->data.string.size = token.data.text.len;
			break;
		case TOK_IDENT:
			ok = snbt_convert_ident(tag, token);
			break;
    }

	return ok;
}

bool snbt_read(tag_t *tag, const stream_t *stream) {
	return snbt_read_value(tag, stream, NULL);
}