#include <pslang/parser/finalize.hpp>
#include <pslang/parser/error.hpp>
#include <pslang/ast/statement.hpp>
#include <pslang/ast/statement_visitor.hpp>

#include <vector>

namespace pslang::parser
{

	namespace
	{

		void validate(ast::statement_list_ptr statements, ast::function_definition * in_function, bool in_loop)
		{
			for (auto & statement : statements->statements)
			{
				if (auto if_chain = std::get_if<ast::if_chain>(statement.get()))
				{
					for (auto const & block : if_chain->blocks)
					{
						validate(block.statements, in_function, in_loop);
					}
				}
				else if (auto while_block = std::get_if<ast::while_block>(statement.get()))
				{
					validate(while_block->statements, in_function, true);
				}
				else if (auto function_definition = std::get_if<ast::function_definition>(statement.get()))
				{
					validate(function_definition->statements, function_definition, in_loop);
				}
				else if (auto return_statement = std::get_if<ast::return_statement>(statement.get()))
				{
					if (!in_function)
						throw parse_error("Return statement outside of function scope", return_statement->location);
					return_statement->node = in_function;
				}
				else if (auto break_statement = std::get_if<ast::break_statement>(statement.get()))
				{
					if (!in_loop)
						throw parse_error("Break without an enclosing loop", break_statement->location);
				}
				else if (auto continue_statement = std::get_if<ast::continue_statement>(statement.get()))
				{
					if (!in_loop)
						throw parse_error("Continue without an enclosing loop", continue_statement->location);
				}
			}
		}

	}

	ast::statement_list_ptr finalize(ast::statement_list_ptr statements)
	{
		validate(statements, nullptr, false);
		return statements;
	}

}