#include <pslang/parser/indented_statement.hpp>
#include <pslang/ast/statement.hpp>

#include <stdexcept>
#include <vector>

namespace pslang::parser
{

	namespace
	{

		ast::statement_list * get_statement_list(ast::expression_ptr &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::assignment &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::variable_declaration &)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::if_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		ast::statement_list * get_statement_list(ast::else_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		ast::statement_list * get_statement_list(ast::else_if_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		// NB: if chain merging happens after retrieving statement list
		ast::statement_list * get_statement_list(ast::if_chain & node)
		{
			return nullptr;
		}

		ast::statement_list * get_statement_list(ast::while_block & node)
		{
			node.statements = std::make_unique<ast::statement_list>();
			return node.statements.get();
		}

		ast::statement_list * get_statement_list(ast::statement & statement)
		{
			return std::visit([](auto & value){ return get_statement_list(value); }, statement);
		}

	}

	ast::statement_list_ptr finilize(indented_statement_list statements)
	{
		ast::statement_list_ptr result = std::make_unique<ast::statement_list>();

		std::vector<ast::statement_list *> stack;
		stack.push_back(result.get());
		std::size_t current_indent = 0;

		for (auto & statement : statements.statements)
		{
			if (statement.indentation > current_indent)
			{
				throw std::runtime_error("Unexpected indent");
			}

			while (statement.indentation < current_indent)
			{
				stack.pop_back();
				--current_indent;
			}

			 // Now statement.indentation == current_indent

			auto list = get_statement_list(*statement.statement);

			if (auto if_block = std::get_if<ast::if_block>(statement.statement.get()))
			{
				ast::if_chain chain;
				chain.blocks.push_back({.condition = std::move(if_block->condition), .statements = std::move(if_block->statements)});
				stack.back()->statements.push_back(std::make_unique<ast::statement>(std::move(chain)));
			}
			else if (auto else_block = std::get_if<ast::else_block>(statement.statement.get()))
			{
				if (stack.back()->statements.empty())
					throw std::runtime_error("Unexpected else block");
				auto chain = std::get_if<ast::if_chain>(stack.back()->statements.back().get());
				if (!chain || chain->blocks.empty() || !chain->blocks.back().condition)
					throw std::runtime_error("Unexpected else block");

				chain->blocks.push_back({.condition = nullptr, .statements = std::move(else_block->statements)});
			}
			else if (auto else_if_block = std::get_if<ast::else_if_block>(statement.statement.get()))
			{
				if (stack.back()->statements.empty())
					throw std::runtime_error("Unexpected else if block");
				auto chain = std::get_if<ast::if_chain>(stack.back()->statements.back().get());
				if (!chain || chain->blocks.empty() || !chain->blocks.back().condition)
					throw std::runtime_error("Unexpected else if block");

				chain->blocks.push_back({.condition = std::move(else_if_block->condition), .statements = std::move(else_if_block->statements)});
			}
			else
			{
				stack.back()->statements.push_back(std::move(statement.statement));
			}

			if (list)
			{
				stack.push_back(list);
				++current_indent;
			}
		}

		return result;
	}

}