#include "compiler.h"
#include "common.h"
#include "emit.h"
#include "lexer.h"

Emitter emitter = {0};
Parser parser = {0};
Arena *arena;

/****************************************************
 *  Scope
 ***************************************************/

void
scope_push()
{
	Scope *child = aalloc(arena, sizeof(Scope));
	child->symbols = List_init(arena);
	child->parent = parser.current_scope;

	parser.current_scope = child;
}

void
scope_pop()
{
	Scope *prev = parser.current_scope->parent;
	parser.current_scope = prev;
}

Symbol *
scope_get_symbol(Scope *scope, const char *name, u32 name_length)
{
	u32 count, i;
	if(!scope) return nil;

	count = scope->symbols->count;
	for(i = 0; i < count; i++) {
		Symbol *sym = List_get(scope->symbols, i);
		if(sleq(sym->name, name, name_length)) return sym;
	}

	return scope_get_symbol(scope->parent, name, name_length);
}

Symbol *
scope_add_symbol(const char *name, u32 name_length, SymbolType type, u32 size)
{
	Symbol *sym = scope_get_symbol(parser.current_scope, name, name_length);
	if(sym != nil) return sym;

	sym = aalloc(arena, sizeof(Symbol));
	scpy(sym->name, name, name_length);
	sym->name_length = name_length;
	sym->type = type;
	sym->secondary_type = SYMBOL_VOID;
	sym->size = size;
	sym->ref = 0;
	if(type == SYMBOL_PLEX || type == SYMBOL_METHOD || type == SYMBOL_TRAIT) {
		sym->args = List_init(arena);
		sym->fields = List_init(arena);
	}

	List_push(arena, parser.current_scope->symbols, sym, sizeof(Symbol));

	return sym;
}

bool
is_global()
{
	return parser.current_scope != nil && parser.current_scope->parent == nil;
}

bool
is_type()
{
	return parser.previous.type >= TOKEN_TYPE_I8 &&
				 parser.previous.type <= TOKEN_TYPE_PTR;
}

/****************************************************
 *  Parser
 ***************************************************/

bool
advance()
{
	parser.previous = parser.current;

	for(;;) {
		parser.current = next_token();
		if(parser.current.type != TOKEN_ERROR) return true;

		return false;
	}
}
bool
check(TokenType type)
{
	return parser.current.type == type;
}

bool
consume(TokenType type)
{
	if(check(type)) {
		advance();
		return true;
	}

	return false;
}

bool
match(TokenType type)
{
	if(!check(type)) return false;
	advance();
	return true;
}

void expression();
void statement();
void declaration();
ParseRule *get_rule(TokenType type);
void parse_precedence(Precedence precedence);

void
binary()
{
	TokenType operatorType = parser.previous.type;
	ParseRule *rule = get_rule(operatorType);

	parse_precedence((Precedence)(rule->precedence + 1));

	switch(operatorType) {
	case TOKEN_BANG_EQ:
		emitter.emit_ne();
		break;
	case TOKEN_EQ_EQ:
		emitter.emit_eq();
		break;
	case TOKEN_GT:
		emitter.emit_gt();
		break;
	case TOKEN_GTE:
		emitter.emit_ge();
		break;
	case TOKEN_LT:
		emitter.emit_lt();
		break;
	case TOKEN_LTE:
		emitter.emit_le();
		break;
	case TOKEN_PLUS:
		emitter.emit_add();
		break;
	case TOKEN_MINUS:
		emitter.emit_sub();
		break;
	case TOKEN_STAR:
		emitter.emit_mul();
		break;
	case TOKEN_SLASH:
		emitter.emit_div();
		break;
	case TOKEN_OPERATOR_AND:
		emitter.emit_and();
		break;
	case TOKEN_OPERATOR_OR:
		emitter.emit_or();
		break;
	case TOKEN_OPERATOR_XOR:
		emitter.emit_xor();
		break;
	case TOKEN_OPERATOR_MOD:
		emitter.emit_mod();
		break;
	case TOKEN_SLL:
		emitter.emit_sll();
		break;
	case TOKEN_SRL:
		emitter.emit_srl();
		break;
	default:
		return;
	}
}

void
literal()
{
	switch(parser.previous.type) {
	case TOKEN_KEYWORD_FALSE:
		emitter.emit_false();
		break;
	case TOKEN_KEYWORD_TRUE:
		emitter.emit_true();
		break;
	case TOKEN_KEYWORD_NIL:
		emitter.emit_nil();
		break;
	default:
		return;
	}
}

void
expression()
{
	parse_precedence(PREC_ASSIGNMENT);
}

void
variable()
{
	Symbol *sym = scope_get_symbol(parser.current_scope, parser.previous.start,
																 parser.previous.length);
	if(sym == nil) {
		emitter.error(parser.previous.start, parser.previous.length,
									parser.previous.line);
	}

	parser.current_type = sym->type;
	emitter.emit_variable(sym, parser.depth);
}

void
cast_type()
{
	SymbolType st = parser.current_type;
	TokenType cast_type = parser.current.type;
	advance();

	switch(st) {
	case SYMBOL_INT: {
		switch(cast_type) {
		case TOKEN_TYPE_NAT: {
			emitter.emit_cast_int_to_nat();
			break;
		}
		case TOKEN_TYPE_REAL: {
			emitter.emit_cast_int_to_real();
			break;
		}
		case TOKEN_TYPE_STR: {
			emitter.emit_cast_int_to_str();
			break;
		}
		default:
			emitter.error("Not castable to this type", 26, parser.previous.line);
		}

		break;
	}
	case SYMBOL_NAT: {
		switch(cast_type) {
		case TOKEN_TYPE_INT: {
			emitter.emit_cast_nat_to_int();
			break;
		}
		case TOKEN_TYPE_REAL: {
			emitter.emit_cast_nat_to_real();
			break;
		}
		case TOKEN_TYPE_STR: {
			emitter.emit_cast_nat_to_str();
			break;
		}
		default:
			emitter.error("Not castable to this type", 26, parser.previous.line);
		}
		break;
	}
	case SYMBOL_REAL: {
		switch(cast_type) {
		case TOKEN_TYPE_NAT: {
			emitter.emit_cast_real_to_nat();
			break;
		}
		case TOKEN_TYPE_INT: {
			emitter.emit_cast_real_to_int();
			break;
		}
		case TOKEN_TYPE_STR: {
			emitter.emit_cast_real_to_str();
			break;
		}
		default:
			emitter.error("Not castable to this type", 26, parser.previous.line);
		}
		break;
	}
	case SYMBOL_STR: {
		switch(cast_type) {
		case TOKEN_TYPE_NAT: {
			emitter.emit_cast_str_to_nat();
			break;
		}
		case TOKEN_TYPE_REAL: {
			emitter.emit_cast_str_to_real();
			break;
		}
		case TOKEN_TYPE_INT: {
			emitter.emit_cast_str_to_int();
			break;
		}
		default:
			emitter.error("Not castable to this type", 26, parser.previous.line);
		}
		break;
	}
	default:
		emitter.error("Not castable to this type", 26, parser.previous.line);
	}
}

void
variable_declaration()
{

	i32 size = 0;
	TokenType tt = parser.previous.type;
	Token var = parser.current;

	switch(tt) {
	case TOKEN_TYPE_INT: {
		size = emitter.emit_int_type(parser.current.start, parser.current.length,
																 parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_INT,
										 size);
		parser.current_type = SYMBOL_INT;
		break;
	}
	case TOKEN_TYPE_NAT: {
		size = emitter.emit_nat_type(parser.current.start, parser.current.length,
																 parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_NAT,
										 size);
		parser.current_type = SYMBOL_NAT;
		break;
	}
	case TOKEN_TYPE_REAL: {
		size = emitter.emit_real_type(parser.current.start, parser.current.length,
																	parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_REAL,
										 size);
		parser.current_type = SYMBOL_REAL;
		break;
	}
	case TOKEN_TYPE_STR: {
		size = emitter.emit_str_type(parser.current.start, parser.current.length,
																 parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_STR,
										 size);
		parser.current_type = SYMBOL_STR;
		break;
	}
	case TOKEN_TYPE_BOOL: {
		size = emitter.emit_bool_type(parser.current.start, parser.current.length,
																	parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_BOOL,
										 size);
		parser.current_type = SYMBOL_BOOL;
		break;
	}
	case TOKEN_TYPE_BYTE: {
		size = emitter.emit_byte_type(parser.current.start, parser.current.length,
																	parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_U8,
										 size);
		parser.current_type = SYMBOL_BYTE;
		break;
	}
	case TOKEN_TYPE_U8: {
		size = emitter.emit_u8_type(parser.current.start, parser.current.length,
																parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_U8,
										 size);
		parser.current_type = SYMBOL_U8;
		break;
	}
	case TOKEN_TYPE_I8: {
		size = emitter.emit_i8_type(parser.current.start, parser.current.length,
																parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_I8,
										 size);
		parser.current_type = SYMBOL_I8;
		break;
	}
	case TOKEN_TYPE_U16: {
		size = emitter.emit_u16_type(parser.current.start, parser.current.length,
																 parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_U16,
										 size);
		parser.current_type = SYMBOL_U16;
		break;
	}
	case TOKEN_TYPE_I16: {
		size = emitter.emit_i16_type(parser.current.start, parser.current.length,
																 parser.depth);
		scope_add_symbol(parser.current.start, parser.current.length, SYMBOL_I16,
										 size);
		parser.current_type = SYMBOL_I16;
		break;
	}
	default: {
		/* probably a plex */
		break;
	}
	}

	advance();

	if(match(TOKEN_EQ)) {
		emitter.emit_set_value();
		expression();
		emitter.emit_constant(var.start, var.length, parser.depth);
	}

	consume(TOKEN_SEMICOLON);
	emitter.emit_end_statement();
	parser.current_type = SYMBOL_UNDEFINED;
}

void
declaration()
{
	/**
	 * maybe a Plex?
	 * lookup plex defs
	 * if not then
	 */
	if(is_type())
		variable_declaration();
	else
		statement();
}

void
print_statement()
{
	expression();
	consume(TOKEN_SEMICOLON);
	emitter.emit_print();
	emitter.emit_end_statement();
}

void
putchar_statement()
{
	expression();
	consume(TOKEN_SEMICOLON);
	emitter.emit_end_statement();
	emitter.emit_putchar();
}

void
getchar_statement()
{
	expression();
	consume(TOKEN_SEMICOLON);
	emitter.emit_end_statement();
	emitter.emit_putchar();
}

void
if_statement()
{
	consume(TOKEN_LPAREN);
	expression();
	consume(TOKEN_RPAREN);
	emitter.emit_jump();
	statement();

	emitter.emit_patch_jump(emitter.ifs);
	if (match(TOKEN_KEYWORD_ELSE)) {
		emitter.else_if_depth++;
		statement();
		emitter.else_if_depth--;
	} 
	if (emitter.else_if_depth == 0) emitter.emit_patch_jump_done(emitter.ifs++);
}

void
block()
{
	while(!check(TOKEN_RBRACE) && !check(TOKEN_EOF)) declaration();

	consume(TOKEN_RBRACE);
}

void
statement()
{
	if(match(TOKEN_LBRACE)) {
		scope_push();
		block();
		scope_pop();
	} else if(match(TOKEN_KEYWORD_IF)) {
		if_statement();
	} else if(match(TOKEN_KEYWORD_PRINT)) {
		print_statement();
	} else if(match(TOKEN_KEYWORD_PUTCHAR)) {
		putchar_statement();
	} else if(match(TOKEN_KEYWORD_GETCHAR)) {
		getchar_statement();
	} else {
		expression();
	}
}

void
grouping()
{
	emitter.emit_open_paren();

	expression();

	emitter.emit_close_paren();
	consume(TOKEN_RPAREN);
}

void
number()
{
	emitter.emit_int(parser.previous.start, parser.previous.length);
}

void
string()
{
	emitter.emit_str(parser.previous.start, parser.previous.length);
}

void
unary()
{
	TokenType operatorType = parser.previous.type;
	parse_precedence(PREC_UNARY);

	switch(operatorType) {
	case TOKEN_MINUS:
		emitter.emit_neg();
		break;
	case TOKEN_BANG:
		emitter.emit_not();
		break;
	default:
		return;
	}
}

ParseRule rules[] = {
		/* TOKEN_ERROR */ {nil, nil, PREC_NONE},
		/* TOKEN_EOF */ {nil, nil, PREC_NONE},
		/* TOKEN_IDENTIFIER */ {variable, nil, PREC_NONE},
		/* TOKEN_LITERAL_INT */ {number, nil, PREC_NONE},
		/* TOKEN_LITERAL_NAT */ {number, nil, PREC_NONE},
		/* TOKEN_LITERAL_REAL */ {number, nil, PREC_NONE},
		/* TOKEN_LITERAL_STR */ {string, nil, PREC_NONE},
		/* TOKEN_TYPE_I8 */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_I16 */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_INT */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_U8 */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_U16 */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_NAT */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_REAL */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_STR */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_BOOL */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_BYTE */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_VOID */ {nil, nil, PREC_NONE},
		/* TOKEN_TYPE_PTR */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_PLEX */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_FN */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_CONST */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_IF */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_IS */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_AS */ {nil, cast_type, PREC_CAST},
		/* TOKEN_KEYWORD_ELSE */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_WHILE */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_FOR */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_RETURN */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_USE */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_INIT */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_THIS */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_OPEN */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_READ */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_WRITE */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_STAT */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_CLOSE */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_LOOP */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_DO */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_NIL */ {literal, nil, PREC_NONE},
		/* TOKEN_KEYWORD_TRUE */ {literal, nil, PREC_NONE},
		/* TOKEN_KEYWORD_FALSE */ {literal, nil, PREC_NONE},
		/* TOKEN_KEYWORD_PUTCHAR */ {nil, nil, PREC_NONE},
		/* TOKEN_KEYWORD_GETCHAR */ {nil, nil, PREC_NONE},
		/* TOKEN_OPERATOR_AND */ {nil, binary, PREC_NONE},
		/* TOKEN_OPERATOR_OR */ {nil, binary, PREC_NONE},
		/* TOKEN_OPERATOR_XOR */ {nil, binary, PREC_NONE},
		/* TOKEN_OPERATOR_MOD */ {nil, binary, PREC_NONE},
		/* TOKEN_BANG */ {unary, nil, PREC_NONE},
		/* TOKEN_BANG_EQ */ {nil, binary, PREC_EQUALITY},
		/* TOKEN_EQ */ {nil, nil, PREC_NONE},
		/* TOKEN_EQ_EQ */ {nil, binary, PREC_EQUALITY},
		/* TOKEN_AND */ {nil, nil, PREC_NONE},
		/* TOKEN_AND_AND */ {nil, nil, PREC_NONE},
		/* TOKEN_PIPE */ {nil, nil, PREC_NONE},
		/* TOKEN_PIPE_PIPE */ {nil, nil, PREC_NONE},
		/* TOKEN_QUESTION */ {nil, nil, PREC_NONE},
		/* TOKEN_QUESTION_DOT */ {nil, nil, PREC_NONE},
		/* TOKEN_PLUS */ {nil, binary, PREC_TERM},
		/* TOKEN_MINUS */ {unary, binary, PREC_TERM},
		/* TOKEN_STAR */ {nil, binary, PREC_FACTOR},
		/* TOKEN_SLASH */ {nil, binary, PREC_FACTOR},
		/* TOKEN_MESH */ {nil, nil, PREC_NONE},
		/* TOKEN_BIG_MONEY */ {nil, nil, PREC_NONE},
		/* TOKEN_GT */ {nil, binary, PREC_COMPARISON},
		/* TOKEN_LT */ {nil, binary, PREC_COMPARISON},
		/* TOKEN_GTE */ {nil, binary, PREC_COMPARISON},
		/* TOKEN_LTE */ {nil, binary, PREC_COMPARISON},
		/* TOKEN_DOT */ {nil, nil, PREC_NONE},
		/* TOKEN_COMMA */ {nil, nil, PREC_NONE},
		/* TOKEN_COLON */ {nil, nil, PREC_NONE},
		/* TOKEN_CARET */ {nil, nil, PREC_NONE},
		/* TOKEN_SEMICOLON */ {nil, nil, PREC_NONE},
		/* TOKEN_LPAREN */ {grouping, nil, PREC_NONE},
		/* TOKEN_RPAREN */ {nil, nil, PREC_NONE},
		/* TOKEN_LBRACE */ {nil, nil, PREC_NONE},
		/* TOKEN_RBRACE */ {nil, nil, PREC_NONE},
		/* TOKEN_LBRACKET */ {nil, nil, PREC_NONE},
		/* TOKEN_RBRACKET */ {nil, nil, PREC_NONE},
		/* TOKEN_ARROW_RIGHT */ {nil, nil, PREC_NONE},
		/* TOKEN_SLL */ {nil, binary, PREC_NONE},
		/* TOKEN_SRL */ {nil, binary, PREC_NONE},
		/*TOKEN_KEYWORD_PRINT*/ {nil, binary, PREC_NONE}};

void
parse_precedence(Precedence precedence)
{
	ParseFn prefixRule;
	ParseFn infixRule;

	advance();
	prefixRule = get_rule(parser.previous.type)->prefix;
	if(prefixRule == nil) return;

	prefixRule();

	while(precedence <= get_rule(parser.current.type)->precedence) {
		advance();
		infixRule = get_rule(parser.previous.type)->infix;
		infixRule();
	}
}

ParseRule *
get_rule(TokenType type)
{
	return &rules[type];
}

bool
compile(Arena *a, Emitter e, char *source)
{
	arena = a;
	emitter = e;
	init_lexer(source);
	advance();
	scope_push();

	emitter.prolog();

	while(!match(TOKEN_EOF)) declaration();

	emitter.epilogue();

	return true;
}