#include <pslang/ast/preprocess.hpp>
#include <pslang/ast/statement_visitor.hpp>
#include <pslang/ast/error.hpp>
#include <pslang/types/type.hpp>

namespace pslang::ast
{

	namespace
	{

		struct validate_visitor
			: const_statement_visitor<validate_visitor>
		{
			using const_statement_visitor::apply;

			void apply(expression_ptr const &) {}

			void apply(assignment const &) {}

			void apply(variable_declaration const &) {}

			void apply(if_chain const & node)
			{
				for (auto const & block : node.blocks)
					apply(*block.statements);
			}

			void apply(while_block const & node)
			{
				apply(*node.statements);
			}

			void apply(break_statement const &)
			{}

			void apply(continue_statement const &)
			{}

			void apply(function_definition const & node)
			{
				apply(*node.statements);
				if (!types::equal(*get_type(*node.return_type), types::unit_type{}))
					if (node.statements->statements.empty() || (true
						&& !std::get_if<return_statement>(node.statements->statements.back().get())
						&& !std::get_if<if_chain>(node.statements->statements.back().get())
						&& !std::get_if<while_block>(node.statements->statements.back().get())
					))
						throw validation_error("Function returning non-unit is missing a return statement in the end", node.location);
			}

			void apply(foreign_function_declaration const &) {}

			void apply(return_statement const &) {}

			void apply(struct_definition const &) {}
		};

	}

	void validate(statement_ptr & root)
	{
		validate_visitor{}.apply(*root);
	}

}