#include "lexer.h"
#include <stdlib.h>
#include <ctype.h>
#include <stdio.h>
#include <string.h>

struct keywords_cell {
	token_type type;
	struct keywords_cell *next;
	char *keyword;
	size_t keyword_length;
};

static struct keywords_cell *keywords[TOKEN_VOLATILE - TOKEN_IF] = {0};

static size_t hash_keyword(char *keyword, size_t length)
{
	size_t sum = 0;
	for (size_t i=0; i < length; i++) {
		sum += (uint8_t) keyword[i];
	}

	return sum % (TOKEN_VOLATILE - TOKEN_IF);
}

static void insert_keyword(char *keyword, size_t length, token_type type)
{
	struct keywords_cell *next = keywords[hash_keyword(keyword, length)];
	if (!next) {
		keywords[hash_keyword(keyword, length)] = calloc(1, sizeof(struct keywords_cell));
		keywords[hash_keyword(keyword, length)]->keyword = keyword;
		keywords[hash_keyword(keyword, length)]->keyword_length = length;
		keywords[hash_keyword(keyword, length)]->type = type;
		keywords[hash_keyword(keyword, length)]->next = NULL;
		return;
	}

	while (next->next) {
		next = next->next;
	}

	next->next = calloc(1, sizeof(struct keywords_cell));
	next->next->keyword = keyword;
	next->next->keyword_length = length;
	next->next->type = type;
	next->next->next = NULL;
}

static token_type get_keyword(char *keyword, size_t length)
{
	struct keywords_cell *cell = keywords[hash_keyword(keyword, length)];
	if (!cell) return 0;
	if (cell->next == NULL) return cell->type;
	while (cell->next != NULL) {
		if (cell->keyword_length != length) {
			cell = cell->next;
			continue;
		}
		if (strncmp(cell->keyword, keyword, length) == 0) {
			return cell->type;
		}
		cell = cell->next;
	}

	if (strncmp(cell->keyword, keyword, length) == 0) {
		return cell->type;
	}

	return 0;
}

static void populate_keywords(void)
{
	insert_keyword("if", 2, TOKEN_IF);
	insert_keyword("else", 4, TOKEN_ELSE);
	insert_keyword("while", 5, TOKEN_WHILE);
	insert_keyword("for", 3, TOKEN_FOR);
	insert_keyword("break", 5, TOKEN_BREAK);
	insert_keyword("case", 4, TOKEN_CASE);
	insert_keyword("const", 5, TOKEN_CONST);
	insert_keyword("do", 2, TOKEN_DO);
	insert_keyword("enum", 4, TOKEN_ENUM);
	insert_keyword("extern", 6, TOKEN_EXTERN);
	insert_keyword("goto", 4, TOKEN_GOTO);
	insert_keyword("inline", 6, TOKEN_INLINE);
	insert_keyword("return", 6, TOKEN_RETURN);
	insert_keyword("signed", 6, TOKEN_SIGNED);
	insert_keyword("unsigned", 8, TOKEN_UNSIGNED);
	insert_keyword("static", 6, TOKEN_STATIC);
	insert_keyword("struct", 6, TOKEN_STRUCT);
	insert_keyword("switch", 6, TOKEN_SWITCH);
	insert_keyword("typedef", 7, TOKEN_TYPEDEF);
	insert_keyword("union", 5, TOKEN_UNION);
	insert_keyword("volatile", 8, TOKEN_VOLATILE);
}

static token *token_create(char *lexeme, size_t length, token_type type)
{
	token *t = calloc(1, sizeof(token));
	t->lexeme = lexeme;
	t->length = length;
	t->type = type;

	return t;
}

static void token_add(token **t, token *new)
{
	if (t) {
		(*t)->next = new;
		new->prev = (*t);
	}

	*t = new;
}

size_t check_symbol(char *src, size_t index, token **last)
{
	switch(src[index]) {
		case '+': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_PLUS_EQ));
				return 2;
			} else if (src[index+1] == '+') {
				token_add(last, token_create(src+index, 2, TOKEN_PLUS_PLUS));
				return 2;

			} else {
				token_add(last, token_create(src+index, 1, TOKEN_PLUS));
				return 1;
			}
		}
		case '-': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_MINUS_EQ));
				return 2;
			} else if (src[index+1] == '-') {
				token_add(last, token_create(src+index, 2, TOKEN_MINUS_MINUS));
				return 2;

			} else if (src[index+1] == '>') {
				token_add(last, token_create(src+index, 2, TOKEN_ARROW));
				return 2;

			} else {
				token_add(last, token_create(src+index, 1, TOKEN_MINUS));
				return 1;
			}
		}
		case '/': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_SLASH_EQ));
				return 2;
			} else {
				token_add(last, token_create(src+index, 1, TOKEN_SLASH));
				return 1;
			}
		}
		case '*': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_STAR_EQ));
				return 2;
			} else {
				token_add(last, token_create(src+index, 1, TOKEN_STAR));
				return 1;
			}
		}
		case '>': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_GREATER_EQ));
				return 2;
			} else if (src[index+1] == '>') {
				if (src[index+2] == '=') {
					token_add(last, token_create(src+index, 2, TOKEN_RSHIFT_EQ));
					return 2;
				} else {
					token_add(last, token_create(src+index, 2, TOKEN_RSHIFT));
					return 2;
				}
			} else {
				token_add(last, token_create(src+index, 1, TOKEN_GREATER));
				return 1;
			}
		}
		case '<': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_LESS_EQ));
				return 2;
			}  else if (src[index+1] == '<') {
				if (src[index+2] == '=') {
					token_add(last, token_create(src+index, 2, TOKEN_LSHIFT_EQ));
					return 2;
				} else {
					token_add(last, token_create(src+index, 2, TOKEN_LSHIFT));
					return 2;
				}
			}  else {
				token_add(last, token_create(src+index, 1, TOKEN_LESS));
				return 1;
			}
		}
		case '!': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_NOT_EQ));
				return 2;
			} else {
				token_add(last, token_create(src+index, 1, TOKEN_NOT));
				return 1;
			}
		}
		case '=': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_EQUAL));
				return 2;
			} else {
				token_add(last, token_create(src+index, 1, TOKEN_ASSIGN));
				return 1;
			}
		}

		case '|': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_BOR_EQ));
				return 2;
			} else if (src[index+1] == '|') {
				token_add(last, token_create(src+index, 2, TOKEN_OR));
				return 2;

			} else {
				token_add(last, token_create(src+index, 1, TOKEN_BOR));
				return 1;
			}
		}
		case '^': {
			if (src[index+1] == '=') {
				token_add(last, token_create(src+index, 2, TOKEN_XOR_EQ));
				return 2;
			} else {
				token_add(last, token_create(src+index, 1, TOKEN_XOR));
				return 1;
			}
		}
		case '{':
			token_add(last, token_create(src+index, 1, TOKEN_LCURLY));
			return 1;
		case '}':
			token_add(last, token_create(src+index, 1, TOKEN_RCURLY));
			return 1;
		case '(':
			token_add(last, token_create(src+index, 1, TOKEN_LPAREN));
			return 1;
		case ')':
			token_add(last, token_create(src+index, 1, TOKEN_RPAREN));
			return 1;
		case '[':
			token_add(last, token_create(src+index, 1, TOKEN_LSQUARE));
			return 1;
		case ']':
			token_add(last, token_create(src+index, 1, TOKEN_RSQUARE));
			return 1;
		case ':':
			token_add(last, token_create(src+index, 1, TOKEN_COLON));
			return 1;
		case ';':
			token_add(last, token_create(src+index, 1, TOKEN_SEMICOLON));
			return 1;
		case ',':
			token_add(last, token_create(src+index, 1, TOKEN_COMMA));
			return 1;
		case '.':
			token_add(last, token_create(src+index, 1, TOKEN_DOT));
			return 1;

		default:
			return 0;
	}
}

token *lexer_parse(char *src)
{
	populate_keywords();
	size_t index = 0;
	token *head = token_create(src, 0, TOKEN_HEAD);
	token *last = head;
	size_t row = 1, col = 1;
	while (src[index] != '\0') {
		if (src[index] == '\n') {
			index++;
			col++;
			row = 1;
			continue;
		}
		if (isspace(src[index])) {
			index++;
			continue;
		}

		size_t skip = check_symbol(src, index, &last);
		if (skip) {
			index += skip;
			continue;
		}

		if (src[index] == '\'') {
			if (src[index+1] != '\\' && src[index+2] != '\'') {
				fprintf(stderr, "%zu:%zu: invalid character\n", row, col);
				return NULL;
			}
			token_add(&last, token_create(src+index+1, 1, TOKEN_CHAR));
			index += 3;
			continue;
		}

		if (src[index] == '\"') {
			index++;
			char *string_start = src+index;
			size_t string_length = index;
			while (src[index++] != '\"') {
				if (src[index] == '\0') {
					fprintf(stderr, "%zu:%zu: unterminated string\n", row, col);
					return NULL;
				}
			}
			string_length = index - string_length - 1;
			token_add(&last, token_create(string_start, string_length, TOKEN_STRING));
			continue;
		}

		if (isalpha(src[index]) || src[index] == '_') {
			char *identifier_start = src+index;
			size_t identifier_length = index;
			while (isalnum(src[index]) || src[index] == '_') {
				index++;
			}
			identifier_length = index - identifier_length;

			token_type keyword = get_keyword(identifier_start, identifier_length);
			if (keyword == 0) {
				token_add(&last, token_create(identifier_start, identifier_length, TOKEN_IDENTIFIER));
			} else {
				token_add(&last, token_create(identifier_start, identifier_length, keyword));
			}
			continue;
		}

		if (isdigit(src[index])) {
			char *number_start = src+index;
			size_t number_length = index;
			uint8_t is_float = 0;

			if (src[index] == '0' && !isdigit(src[index+1]) && src[index+1] != '.' && (isdigit(src[index+2]) || ((src[index+2] >= 65 && src[index+2] <= 70) && (src[index+2] >= 97 && src[index+2] <= 102)))) {
				if (src[index+1] == 'x') {
					index += 2;
				} else if (src[index+1] == 'o') {
					index += 2;
				} else if (src[index+1] == 'b') {
					index += 2;
				} else {
					fprintf(stderr, "%zu:%zu: invalid integer\n", row, col);
					return NULL;
				}
			}
			index++;
			while (isdigit(src[index]) || src[index] == '.') {
				if (src[index] == '.') {
					if (is_float) {
						fprintf(stderr, "%zu:%zu: invalid float\n", row, col);
						return NULL;
					}
					is_float = 1;
				}
				index++;
			}


			number_length = index - number_length;
			token_add(&last, token_create(number_start, number_length, is_float ? TOKEN_FLOAT : TOKEN_INTEGER));
			continue;
		}
		index ++;
	}

	return head;
}