#include "parser.hpp"

#include "bytebuf.hpp"

// Example:
// 	if chance(45) and seesenemy then
// 		setstate(@wiz_shoot0)
// 		stopmoving
// 	else
// 		move(0, 1.6)
// 	endif
// Parses as:
// 	IFELSE
// 		BINARYOP	and
// 			FUNC
// 				NUMLIT 45
// 			FUNC
//		SEQ
//			FUNC setstate
//				REFLIT wiz_shoot0
//			FUNC stopmoving
//		SEQ
//			FUNC move
//				NUMLIT 0
//				NUMLIT 1.6

// Defines when to stop parsing statements.
// If a non-matching end type is found, it will be treated as an error.
enum SeqEndType {
	SEQEND_BLOCKEND,
	SEQEND_ELSE,	// 'endif' is also valid here
	SEQEND_ENDIF,
};
ASTNode ParseSequence(TokenReader &reader, SeqEndType endType);

ASTNode ParseIfElse(TokenReader &reader);
ASTNode ParseFunctionCall(TokenReader &reader);

// Similar to ParseSequence, but only one statement is parsed,
// and the resulting code must push to the stack 
ASTNode ParseExpression(TokenReader &reader);


#define NUM_KEYWORDS 7
static const char *keywords[NUM_KEYWORDS] = {
	"if", "then", "else", "endif", "not", "and", "or",
};

bool IsKeyword(const std::string &);

ASTNode ParseScript(TokenReader &reader)
{
	reader.ReadToken(TOK_BLOCK_BEGIN);
	ASTNode ast = ParseSequence(reader, SEQEND_BLOCKEND);
	reader.ReadToken(TOK_BLOCK_END);
	return ast;
}

void DumpAST(const ASTNode &root, int level)
{
	for (int i = 0; i < level; i++) {
		putchar('\t');
	}
	if (root.value.contents.empty())
		printf("%d\n", root.type);
	else
		printf("%d -> %s\n", root.type, root.value.contents.c_str());
	for (const auto &node : root.subNodes) {
		DumpAST(node, level + 1);
	}
}

// An expression is a statement that results in a value being pushed to the stack.
// Note that values should be pushed to the stack in the opposite order that they appear in the code.

// shoot(@fireball, $dmg + $atk, $spd + ($int + 4), 0)
//
// ;; 0
// push 0
// ;; $spd + ($int + 4)
// push 4
// get int
// add			; ($int + 4)
// get spd
// add

bool IsKeyword(const std::string &str)
{
	for (int i = 0; i < NUM_KEYWORDS; i++)
	{
		if (str == keywords[i])
			return true;
	}
	return false;
}

ASTNode ParseSequence(TokenReader &reader, SeqEndType endType)
{
	ASTNode root;
	root.type = ASTNODE_SEQ;
	while (reader.PeekToken().type != TOK_BLOCK_END) {
		token tok = reader.PeekToken();
		// Each line of the program can be an if statement, a function call,
		// or an assignment to a register.
		if (tok.type == TOK_NAME) {
			if (IsKeyword(tok.contents)) {
				if (tok.contents == "if") {
					root.subNodes.push_back(ParseIfElse(reader));
				} else if (tok.contents == "else") {
					if (endType != SEQEND_ELSE)
						PrintErrorMsg(tok.filepos, ERR_LVL_ERROR, "Unexpected 'else'");
					break;
				} else if (tok.contents == "endif") {
					if (endType != SEQEND_ENDIF && endType != SEQEND_ELSE)
						PrintErrorMsg(tok.filepos, ERR_LVL_ERROR, "Unexpected 'endif'");
					break;
				}
			} else {
				// Unquoted string that is not a keyword, treat it as a function.
				root.subNodes.push_back(ParseFunctionCall(reader));
			}
		} else if (tok.type == TOK_REGISTER) {
			// Must have an assignment to a register in this case.
			ASTNode assignNode, registerNode;

			registerNode.type = ASTNODE_LVALUE;
			registerNode.value = reader.ReadToken();

			token optok = reader.ReadToken();
			if (optok.type != TOK_ASSIGN && optok.type != TOK_PLUSEQU && optok.type != TOK_MINUSEQU) {
				PrintErrorMsg(optok.filepos, ERR_LVL_ERROR, "Expected assignment to register");
				return root;
			}

			assignNode.type = ASTNODE_ASSIGNMENT;
			assignNode.value = optok;
			assignNode.subNodes.push_back(registerNode);
			assignNode.subNodes.push_back(ParseExpression(reader));

			root.subNodes.push_back(assignNode);
		} else {
			PrintErrorMsg(tok.filepos, ERR_LVL_ERROR, "Invalid statement");
			return root;
		}
	}
	return root;
}

ASTNode ParseIfElse(TokenReader &reader)
{
	ASTNode ifNode;
	ifNode.type = ASTNODE_IFELSE;

	// Consume the 'if' token
	token ifToken = reader.ReadToken(TOK_NAME);

	ifNode.subNodes.push_back(ParseExpression(reader));

	// Consume the 'then' token, this could probably be removed.
	token thenToken = reader.ReadToken();
	if (thenToken.type != TOK_NAME || thenToken.contents != "then") {
		PrintErrorMsg(ifToken.filepos, ERR_LVL_ERROR, "Expected 'then' after 'if' statement");
		return ifNode;
	}

	ifNode.subNodes.push_back(ParseSequence(reader, SEQEND_ELSE));

	token elseOrEndif = reader.ReadToken();
	if (elseOrEndif.type != TOK_NAME || (elseOrEndif.contents != "else" && elseOrEndif.contents != "endif")) {
		PrintErrorMsg(ifToken.filepos, ERR_LVL_ERROR, "Expected 'else' or 'endif' within 'if' statement");
		return ifNode;
	}

	if (elseOrEndif.contents == "else") {
		ifNode.subNodes.push_back(ParseSequence(reader, SEQEND_ENDIF));
		token endifToken = reader.ReadToken();
		if (endifToken.type != TOK_NAME || endifToken.contents != "endif") {
			PrintErrorMsg(ifToken.filepos, ERR_LVL_ERROR, "Expected 'endif' at end of if statement");
			return ifNode;
		}
	}

	return ifNode;
}

ASTNode ParseFunctionCall(TokenReader &reader)
{
	ASTNode funcNode;
	funcNode.type = ASTNODE_FUNC;
	funcNode.value = reader.ReadToken(TOK_NAME);

	if (reader.PeekToken().type == TOK_PARAMS_BEGIN) {
		reader.ReadToken(TOK_PARAMS_BEGIN);

		// Parse each function parameter as an expression
		do {
			funcNode.subNodes.push_back(ParseExpression(reader));
			if (reader.PeekToken().type == TOK_PARAMS_END)
				break;
			reader.ReadToken(TOK_PARAMS_SEPARATOR);
		} while (!reader.EndOfFile());
		reader.ReadToken(TOK_PARAMS_END);
	}

	return funcNode;
}

bool IsTokenBinaryOp(const token &tok)
{
	if (tok.type == TOK_PLUS || tok.type == TOK_MINUS
		|| tok.type == TOK_COMPARE
		|| tok.type == TOK_LESSTHAN || tok.type == TOK_GREATERTHAN
		|| tok.type == TOK_LESSEQU || tok.type == TOK_GREATEREQU)
	{
		return true;
	}
	if (tok.type == TOK_NAME
		&& (tok.contents == "and" || tok.contents == "or"))
	{
		return true;
	}
	return false;
}

// -4 + rand(2, 5 - $reg) - 3
//	BINOP	+
//		UNOP -
//			LIT 4
//		BINOP	-
//			FUNC rand
//				LIT 2
//				BINOP	-
//					LIT 5
//					LIT $ref
//			LIT 3
ASTNode ParseExpression(TokenReader &reader)
{
	// An expression can be a:
	//  - Function call
	//  - Register
	//  - Numeric literal
	//	- Unary or binary operator

	// This doesn't handle order of operations at the moment,
	// since only addition and subtraction have been implemented so far. 

	ASTNode result;
	
	token tok = reader.PeekToken();
	switch (tok.type)
	{
	case TOK_NAME:
		if (tok.contents == "not") {
			result.type = ASTNODE_UNARYOP;
			result.value = reader.ReadToken();
			result.subNodes.push_back(ParseExpression(reader));
		} else {
			result = ParseFunctionCall(reader);
		}
		break;
	case TOK_REGISTER:
		result.type = ASTNODE_REGISTER;
		result.value = reader.ReadToken();
		break;
	case TOK_STRING:
	case TOK_REFERENCE:
	case TOK_NUMBER:
		result.type = ASTNODE_LITERAL;
		result.value = reader.ReadToken();
		break;
	case TOK_MINUS:
		result.type = ASTNODE_UNARYOP;
		result.value = reader.ReadToken();
		result.subNodes.push_back(ParseExpression(reader));
		break;
	default:
		break;
	}

	// If the next token is a binary operator, keep reading the expression.
	// Otherwise, we have reached the end.
	if (IsTokenBinaryOp(reader.PeekToken())) {
		ASTNode operatorNode;
		operatorNode.type = ASTNODE_BINARYOP;
		operatorNode.value = reader.ReadToken();
		operatorNode.subNodes.push_back(result);
		operatorNode.subNodes.push_back(ParseExpression(reader));
		return operatorNode;
	} else {
		return result;
	}
}
