#include "../parser.hpp"
#include <algorithm>

std::string LocationText(const Token& t)
{
	return "at line " + std::to_string(t.line) + " column " + std::to_string(t.column);
}

Token& Parser::Current()
{
	if (_pos < _tokens.size())
		return _tokens[_pos];

	throw std::runtime_error("Unidentified token!");
}

Token& Parser::Consume()
{
	if (_pos < _tokens.size())
		return _tokens[_pos++];

	throw std::runtime_error("Unidentified token!");
}

Token& Parser::Peek(size_t offset)
{
	if (_pos + offset < _tokens.size())
		return _tokens[_pos + offset];

	throw std::runtime_error("Unidentified token!");
}

bool Parser::Match(TokenType type, const std::string& value)
{
	const auto& current = Current();

	if (current.type != type || current.value != value)
		return false;

	Consume();
	return true;
}

Token& Parser::ShouldBe(TokenType type, const std::string& value)
{
	const auto& current = Current();

	if (current.type != type)
		throw std::runtime_error("Syntax Error: Expected token type '" + std::string(TokenTypeToString(type)) +
			"' but got '" + std::string(TokenTypeToString(current.type)) +
			"' (value: '" + current.value + "') " + LocationText(current));

	if (!value.empty() && value != current.value)
		throw std::runtime_error("Syntax Error: Expected specific value '" + value +
			"' but got '" + current.value + "' " + LocationText(current));

	return Consume();
}

bool Parser::Match(TokenType type, char value)
{
	return Match(type, std::string(1, value));
}

Token& Parser::ShouldBe(TokenType type, char value)
{
	return ShouldBe(type, std::string(1, value));
}

Parser::Parser(const std::vector<Token>& tokens) : _tokens(tokens) { }

Program Parser::ParseProgram()
{
	Program program;

	while (_pos < _tokens.size())
		program.statements.push_back(ParseStatement());

	return program;
}

using namespace Dlang::Language;
#pragma region STATEMENTS

std::unique_ptr<Statement> Parser::ParseStatement()
{
	if (Match(TokenType::KEYWORD, kw_var))
		return ParseVariableDeclaration();

	if (Match(TokenType::KEYWORD, kw_while))
		return ParseWhileDeclaration();

	if (Match(TokenType::KEYWORD, kw_for))
		return ParseForDeclaration();

	if (Match(TokenType::KEYWORD, kw_foreach))
		return ParseForeachDeclaration();

	if (Match(TokenType::KEYWORD, kw_if))
		return ParseIfDeclaration();

	if (Match(TokenType::KEYWORD, kw_func))
		return ParseFunctionDeclaration();

	if (Match(TokenType::KEYWORD, kw_return))
		return ParseReturnStatement();

	if (Match(TokenType::KEYWORD, kw_include))
		return ParseIncludeStatement();

	auto expr = ParseExpression();

	if (Match(TokenType::SYNTAX, op_equals))
	{
		auto op = Peek(-1).value;
		if (!dynamic_cast<IdentifierExpression*>(expr.get()) &&
			!dynamic_cast<MemberAccessExpression*>(expr.get()) &&
			!dynamic_cast<IndexAccessExpression*>(expr.get()))
		{
			throw std::runtime_error("Invalid left-hand side in assignment");
		}

		auto value = ParseExpression();
		ShouldBe(TokenType::SYNTAX, op_endl);

		auto assignment = std::make_unique<Assignment>();
		assignment->lhs = std::move(expr);
		assignment->value = std::move(value);
		assignment->assignmentOperator = op;
		return assignment;
	}
	
	if (Match(TokenType::SYNTAX, op_add_assign) || Match(TokenType::SYNTAX, op_sub_assign) || Match(TokenType::SYNTAX, op_mul_assign) || Match(TokenType::SYNTAX, op_div_assign))
	{
		std::string op = Peek(-1).value;

		if (!dynamic_cast<IdentifierExpression*>(expr.get()) &&
			!dynamic_cast<MemberAccessExpression*>(expr.get()) &&
			!dynamic_cast<IndexAccessExpression*>(expr.get()))
		{
			throw std::runtime_error("Invalid left-hand side in assignment");
		}

		auto value = ParseExpression();
		ShouldBe(TokenType::SYNTAX, op_endl);

		auto assignment = std::make_unique<Assignment>();
		assignment->lhs = std::move(expr);
		assignment->value = std::move(value);
		assignment->assignmentOperator = op;
		return assignment;
	}


	ShouldBe(TokenType::SYNTAX, op_endl);

	auto stmt = std::make_unique<ExpressionStatement>();
	stmt->expression = std::move(expr);
	return stmt;
}

std::unique_ptr<Statement> Parser::ParseVariableDeclaration()
{
	auto name = ShouldBe(TokenType::IDENTIFIER);
	ShouldBe(TokenType::SYNTAX, op_equals);
	auto initializer = ParseExpression();
	ShouldBe(TokenType::SYNTAX, op_endl);

	auto declaration = std::make_unique<VariableDeclaration>();
	declaration->name = name.value;
	declaration->initializer = std::move(initializer);

	return declaration;
}

std::unique_ptr<Statement> Parser::ParseAssignment()
{
	auto lhs = ParseExpression();

	if (!dynamic_cast<IdentifierExpression*>(lhs.get()) &&
		!dynamic_cast<MemberAccessExpression*>(lhs.get()))
	{
		throw std::runtime_error("Invalid left-hand side in assignment");
	}

	ShouldBe(TokenType::SYNTAX, op_equals);
	auto value = ParseExpression();
	ShouldBe(TokenType::SYNTAX, op_endl);

	auto assignment = std::make_unique<Assignment>();
	assignment->lhs = std::move(lhs);
	assignment->value = std::move(value);

	return assignment;
}

std::unique_ptr<Statement> Parser::ParseWhileDeclaration()
{
	ShouldBe(TokenType::SYNTAX, op_exprstart);

	auto stmt = std::make_unique<WhileStatement>();
	stmt->condition = ParseExpression();
	
	ShouldBe(TokenType::SYNTAX, op_exprend);
	ShouldBe(TokenType::KEYWORD, kw_do);

	while (!Match(TokenType::KEYWORD, kw_end))
		stmt->body.push_back(ParseStatement());

	return stmt;
}

std::unique_ptr<Statement> Parser::ParseForDeclaration()
{
	ShouldBe(TokenType::SYNTAX, op_exprstart);
	auto varName = ShouldBe(TokenType::IDENTIFIER);
	ShouldBe(TokenType::SYNTAX, op_separator);
	auto from = ShouldBe(TokenType::NUMBER);
	ShouldBe(TokenType::SYNTAX, op_separator);
	auto to = ShouldBe(TokenType::NUMBER);

	auto stmt = std::make_unique<ForStatement>();
	stmt->var = varName.value;
	stmt->from = std::stod(from.value);
	stmt->to = std::stod(to.value);

	if (Match(TokenType::SYNTAX, op_separator))
	{
		auto step = ShouldBe(TokenType::NUMBER);
		stmt->step = std::stod(step.value);
	}

	ShouldBe(TokenType::SYNTAX, op_exprend);
	ShouldBe(TokenType::KEYWORD, kw_do);

	while (!Match(TokenType::KEYWORD, kw_end))
		stmt->body.push_back(ParseStatement());

	return stmt;
}

std::unique_ptr<Statement> Parser::ParseForeachDeclaration()
{
	ShouldBe(TokenType::SYNTAX, op_exprstart);
	auto varName = ShouldBe(TokenType::IDENTIFIER);
	if (Match(TokenType::KEYWORD, kw_in))
	{
		auto list = ShouldBe(TokenType::IDENTIFIER);
		ShouldBe(TokenType::SYNTAX, op_exprend);
		ShouldBe(TokenType::KEYWORD, kw_do);

		auto stmt = std::make_unique<ForeachStatement>();
		stmt->var = varName.value;
		stmt->list = list.value;

		while (!Match(TokenType::KEYWORD, kw_end))
			stmt->body.push_back(ParseStatement());

		return stmt;
	}
	else {
		ShouldBe(TokenType::SYNTAX, op_separator);
		auto value = ShouldBe(TokenType::IDENTIFIER);
		ShouldBe(TokenType::KEYWORD, kw_in);
		auto table = ShouldBe(TokenType::IDENTIFIER);
		ShouldBe(TokenType::SYNTAX, op_exprend);
		ShouldBe(TokenType::KEYWORD, kw_do);

		auto stmt = std::make_unique<ForeachTableStatement>();
		stmt->varKey = varName.value;
		stmt->varValue = value.value;
		stmt->table = table.value;

		while (!Match(TokenType::KEYWORD, kw_end))
			stmt->body.push_back(ParseStatement());

		return stmt;
	}
}

std::unique_ptr<Statement> Parser::ParseIfDeclaration()
{
	ShouldBe(TokenType::SYNTAX, op_exprstart);
	auto expr = ParseExpression();
	ShouldBe(TokenType::SYNTAX, op_exprend);
	ShouldBe(TokenType::KEYWORD, kw_then);

	auto stmt = std::make_unique<IfStatement>();
	stmt->expression = std::move(expr);

	Token t;
	while (!Match(TokenType::KEYWORD, kw_else) && !Match(TokenType::KEYWORD, kw_end))
		stmt->body.push_back(ParseStatement());

	if (Peek(-1).type == TokenType::KEYWORD && Peek(-1).value == kw_else)
	{
		if (Match(TokenType::KEYWORD, kw_if))
			stmt->elseBody.push_back(ParseIfDeclaration());
		else
			while (!Match(TokenType::KEYWORD, kw_end))
				stmt->elseBody.push_back(ParseStatement());
	}

	return stmt;
}

std::unique_ptr<Statement> Parser::ParseFunctionDeclaration()
{
	auto funcName = ShouldBe(TokenType::IDENTIFIER);
	ShouldBe(TokenType::SYNTAX, op_exprstart);
	std::vector<std::string> params;
	if (!Match(TokenType::SYNTAX, op_exprend))
	{
		do
		{
			Token param = ShouldBe(TokenType::IDENTIFIER);
			params.push_back(param.value);
		} while (Match(TokenType::SYNTAX, op_separator));

		ShouldBe(TokenType::SYNTAX, op_exprend);
	}

	auto stmt = std::make_unique<FunctionStatement>();
	stmt->name = funcName.value;
	stmt->parameters = params;
	
	while (!Match(TokenType::KEYWORD, kw_end))
		stmt->body.push_back(ParseStatement());

	return stmt;
}

std::unique_ptr<Statement> Parser::ParseReturnStatement()
{
	auto stmt = std::make_unique<ReturnStatement>();

	if (!Match(TokenType::SYNTAX, op_endl))
	{
		stmt->expression = ParseExpression();
		ShouldBe(TokenType::SYNTAX, op_endl);
	}

	return stmt;
}

std::unique_ptr<Statement> Parser::ParseIncludeStatement()
{
	auto stmt = std::make_unique<IncludeStatement>();

	ShouldBe(TokenType::SYNTAX, op_exprstart);
	auto path = ShouldBe(TokenType::STRING).value;
	ShouldBe(TokenType::SYNTAX, op_exprend);
	ShouldBe(TokenType::SYNTAX, op_endl);
	stmt->path = path;

	return stmt;
}

#pragma endregion

#pragma region EXPRESSIONS

std::unique_ptr<Expression> Parser::ParseFunctionCall(std::unique_ptr<Expression> callee)
{
	auto stmt = std::make_unique<CallExpression>();
	stmt->callee = std::move(callee);

	if (!Match(TokenType::SYNTAX, op_exprend))
	{
		do {
			stmt->expressions.push_back(ParseExpression());
		} while (Match(TokenType::SYNTAX, op_separator));

		if (!Match(TokenType::SYNTAX, op_exprend))
			throw std::runtime_error("Expected ')' after arguments " + LocationText(Current()));
	}

	return stmt;
}

std::unique_ptr<Expression> Parser::ParseExpression()
{
	return ParseLogicalAndOr();
}

std::unique_ptr<Expression> Parser::ParsePrimary()
{
	if (Match(TokenType::SYNTAX, op_exprstart))
	{
		auto expression = ParseExpression();

		if (!Match(TokenType::SYNTAX, op_exprend))
			throw std::runtime_error("Expected ')' " + LocationText(Current()));

		return expression;
	}

	if (Match(TokenType::SYNTAX, op_arraystart))
		return ParseTableLiteral();

	if (Match(TokenType::SYNTAX, op_liststart))
		return ParseListLiteral();

	if (Match(TokenType::KEYWORD, kw_func))
		return ParseFunctionExpression();

	if (Match(TokenType::SYNTAX, op_size))
		return ParseSizeExpression();

	const auto& current = Current();

	if (current.type == TokenType::NUMBER)
	{
		auto token = ShouldBe(TokenType::NUMBER);
		
		auto expression = std::make_unique<NumberExpression>();
		expression->value = std::stod(token.value);

		return expression;
	}

	if (current.type == TokenType::BOOLEAN)
	{
		auto token = ShouldBe(TokenType::BOOLEAN);

		auto expression = std::make_unique<BooleanExpression>();
		expression->value = (token.value == kw_true ? true : false);

		return expression;
	}

	if (current.type == TokenType::STRING)
	{
		auto token = ShouldBe(TokenType::STRING);

		auto expression = std::make_unique<StringExpression>();
		expression->value = token.value;

		return expression;
	}

	if (current.type == TokenType::IDENTIFIER)
	{
		auto token = ShouldBe(TokenType::IDENTIFIER);

		auto idExpr = std::make_unique<IdentifierExpression>();
		idExpr->name = token.value;

		std::unique_ptr<Expression> left;

		if (Match(TokenType::SYNTAX, op_exprstart))
			left = ParseFunctionCall(std::move(idExpr));
		else
			left = std::move(idExpr);

		while (Match(TokenType::SYNTAX, op_dot) || Match(TokenType::SYNTAX, op_liststart))
		{
			if (Peek(-1).value[0] == op_dot)
			{
				auto memberExpr = std::make_unique<MemberAccessExpression>();
				memberExpr->object = std::move(left);
				memberExpr->memberName = ShouldBe(TokenType::IDENTIFIER).value;
				left = std::move(memberExpr);

				if (Match(TokenType::SYNTAX, op_exprstart))
				{
					left = ParseFunctionCall(std::move(left));
				}
			}
			else
			{
				auto indexExpr = ParseExpression();

				ShouldBe(TokenType::SYNTAX, op_listend);

				auto indexAccess = std::make_unique<IndexAccessExpression>();
				indexAccess->object = std::move(left);
				indexAccess->memberIndex = std::move(indexExpr);

				left = std::move(indexAccess);
			}
		}

		return left;
	}

	std::string tokenDesc = std::string(TokenTypeToString(current.type));
	if (!current.value.empty()) {
		tokenDesc += " ('" + current.value + "')";
	}
	throw std::runtime_error("Expected expression, but found unexpected token type " + tokenDesc + " " + LocationText(current));
}

std::unique_ptr<Expression> Parser::ParseMultiplicative()
{
	auto left = ParsePrimary();

	while (Current().type == TokenType::SYNTAX &&
		(Current().value[0] == op_multi || Current().value[0] == op_div))
	{
		auto op = Consume();
		auto right = ParsePrimary();

		auto expression = std::make_unique<BinaryExpression>();
		expression->left = std::move(left);
		expression->op = op.value;
		expression->right = std::move(right);

		left = std::move(expression);
	}

	return left;
}

std::unique_ptr<Expression> Parser::ParseAdditive()
{
	auto left = ParseMultiplicative();

	while (Current().type == TokenType::SYNTAX &&
		(Current().value[0] == op_add || Current().value[0] == op_sub))
	{
		auto op = Consume();

		auto right = ParseMultiplicative();

		auto expression = std::make_unique<BinaryExpression>();

		expression->left = std::move(left);
		expression->op = op.value;
		expression->right = std::move(right);

		left = std::move(expression);
	}

	return left;
}

std::unique_ptr<Expression> Parser::ParseComparison()
{
	auto left = ParseAdditive();

	while (
		Match(TokenType::SYNTAX, op_above) ||
		Match(TokenType::SYNTAX, op_below) ||
		Match(TokenType::SYNTAX, op_eq) ||
		Match(TokenType::SYNTAX, op_gte) ||
		Match(TokenType::SYNTAX, op_lte) ||
		Match(TokenType::SYNTAX, op_neq)
	) {
		auto op = Peek(-1);
		auto right = ParseAdditive();

		auto expr = std::make_unique<BinaryExpression>();
		expr->left = std::move(left);
		expr->op = op.value;
		expr->right = std::move(right);

		left = std::move(expr);
	}

	return left;
}

std::unique_ptr<Expression> Parser::ParseTableLiteral()
{
	auto expr = std::make_unique<TableLiteralExpression>();
	
	while (!Match(TokenType::SYNTAX, op_arrayend))
	{
		std::string key = ShouldBe(TokenType::IDENTIFIER).value;
		ShouldBe(TokenType::SYNTAX, op_colon);

		auto valueExpr = ParseExpression();
		expr->fields.push_back({ key, std::move(valueExpr) });

		Match(TokenType::SYNTAX, op_separator);
	}

	return expr;
}

std::unique_ptr<Expression> Parser::ParseListLiteral()
{
	auto expr = std::make_unique<ListLiteralExpression>();

	while (!Match(TokenType::SYNTAX, op_listend))
	{
		auto valueExpr = ParseExpression();
		expr->fields.push_back(std::move(valueExpr));

		Match(TokenType::SYNTAX, op_separator);
	}

	return expr;
}

std::unique_ptr<Expression> Parser::ParseLogicalAndOr()
{
	auto left = ParseComparison();

	while (Match(TokenType::SYNTAX, op_and) || Match(TokenType::SYNTAX, op_or))
	{
		auto op = Peek(-1);
		auto right = ParseComparison();

		auto expr = std::make_unique<BinaryExpression>();
		expr->left = std::move(left);
		expr->op = op.value;
		expr->right = std::move(right);

		left = std::move(expr);
	}

	return left;
}

std::unique_ptr<Expression> Parser::ParseFunctionExpression()
{
	ShouldBe(TokenType::SYNTAX, op_exprstart);
	std::vector<std::string> params;
	if (!Match(TokenType::SYNTAX, op_exprend))
	{
		do
		{
			Token param = ShouldBe(TokenType::IDENTIFIER);
			params.push_back(param.value);
		} while (Match(TokenType::SYNTAX, op_separator));

		ShouldBe(TokenType::SYNTAX, op_exprend);
	}

	auto expr = std::make_unique<FunctionExpression>();
	expr->parameters = params;

	while (!Match(TokenType::KEYWORD, kw_end))
		expr->body.push_back(ParseStatement());

	return expr;
}

std::unique_ptr<Expression> Parser::ParseSizeExpression()
{
	auto target = ShouldBe(TokenType::IDENTIFIER).value;

	auto stmt = std::make_unique<SizeExpression>();
	stmt->identifier = target;

	return stmt;
}

#pragma endregion