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

#include <unordered_set>
#include <vector>

namespace pslang::ast
{

	namespace
	{

		struct scope
		{
			std::unordered_set<std::string> functions;
			std::unordered_set<std::string> structs;
			std::unordered_set<std::string> variables;

			bool contains_transitive(std::string const & name) const
			{
				return false
					|| (functions.count(name) > 0)
					|| (structs.count(name) > 0)
					;
			}

			bool contains_type(std::string const & name) const
			{
				return false
					|| (structs.count(name) > 0)
					;
			}

			bool contains(std::string const & name) const
			{
				return false
					|| (functions.count(name) > 0)
					|| (structs.count(name) > 0)
					|| (variables.count(name) > 0)
					;
			}

			bool contains(std::string const & name, bool crossed_function_scope) const
			{
				if (crossed_function_scope)
					return contains_transitive(name);
				return contains(name);
			}

			bool is_function_scope = false;
			bool is_global_scope = false;
		};

		struct resolve_identifiers_visitor
			: type_visitor<resolve_identifiers_visitor>
			, expression_visitor<resolve_identifiers_visitor>
			, statement_visitor<resolve_identifiers_visitor>
		{
			std::vector<scope> scopes;

			using type_visitor::apply;
			using expression_visitor::apply;
			using statement_visitor::apply;

			void apply(types::unit_type const &)
			{}

			void apply(types::primitive_type const &)
			{}

			void apply(array_type const & array_type)
			{
				apply(*array_type.element_type);
			}

			void apply(function_type const & function_type)
			{
				for (auto const & argument : function_type.arguments)
					apply(*argument);
				apply(*function_type.result);
			}

			void apply(type_identifier & identifier)
			{
				for (auto it = scopes.rbegin(); it != scopes.rend(); ++it)
				{
					if (it->contains_type(identifier.name))
					{
						identifier.level = it.base() - scopes.begin() - 1;
						return;
					}
				}

				throw parse_error("Identifier \"" + identifier.name + "\" not found", identifier.location);
			}

			void apply(literal const &)
			{}

			void apply(identifier & identifier)
			{
				if (types::builtin_type(identifier.name))
					return;

				bool crossed_function_scope = false;
				for (auto it = scopes.rbegin(); it != scopes.rend(); ++it)
				{
					if (it->contains(identifier.name, crossed_function_scope && !it->is_global_scope))
					{
						identifier.level = it.base() - scopes.begin() - 1;
						return;
					}

					crossed_function_scope |= it->is_function_scope;
				}

				throw parse_error("Identifier \"" + identifier.name + "\" not found", identifier.location);
			}

			void apply(unary_operation const & unary_operation)
			{
				apply(*unary_operation.arg1);
			}

			void apply(binary_operation const & binary_operation)
			{
				apply(*binary_operation.arg1);
				apply(*binary_operation.arg2);
			}

			void apply(cast_operation const & cast_operation)
			{
				apply(*cast_operation.expression);
				apply(*cast_operation.type);
			}

			void apply(function_call & function_call)
			{
				if (function_call.function)
					apply(*function_call.function);
				if (function_call.type)
					apply(*function_call.type);
				for (auto const & argument : function_call.arguments)
					apply(*argument);

				if (auto id = std::get_if<identifier>(function_call.function.get()))
				{
					if (auto type = types::builtin_type(id->name))
					{
						if (auto unit_type = std::get_if<types::unit_type>(type.get()))
							function_call.type = std::make_unique<ast::type>(*unit_type);
						else if (auto primitive_type = std::get_if<types::primitive_type>(type.get()))
							function_call.type = std::make_unique<ast::type>(*primitive_type);
						else
							throw invalid_ast_error("Unknown built-in type \"" + id->name + "\"", get_location(*function_call.function));

						function_call.function = nullptr;
					}
					else if (scopes.at(id->level).structs.contains(id->name))
					{
						function_call.type = std::make_unique<ast::type>(type_identifier{.name = id->name, .location = id->location, .level = id->level});
						function_call.function = nullptr;
					}
				}
			}

			void apply(array const & array)
			{
				for (auto const & element : array.elements)
					apply(*element);
			}

			void apply(array_access const & array_access)
			{
				apply(*array_access.array);
				apply(*array_access.index);
			}

			void apply(field_access const & field_access)
			{
				apply(*field_access.object);
			}

			void apply(expression_ptr const & expression_ptr)
			{
				apply(*expression_ptr);
			}

			void apply(assignment const & assignment)
			{
				apply(assignment.lhs);
				apply(assignment.rhs);
			}

			void apply(variable_declaration const & variable_declaration)
			{
				if (scopes.back().contains(variable_declaration.name))
					throw parse_error("Identifier \"" + variable_declaration.name + "\" is already defined at this scope", variable_declaration.location);

				if (variable_declaration.type)
					apply(*variable_declaration.type);
				apply(*variable_declaration.initializer);

				scopes.back().variables.insert(variable_declaration.name);
			}

			void apply(if_block const & if_block)
			{
				throw invalid_ast_error("if blocks cannot be present in the final AST", if_block.location);
			}

			void apply(else_if_block const & else_if_block)
			{
				throw invalid_ast_error("else if blocks cannot be present in the final AST", else_if_block.location);
			}

			void apply(else_block const & else_block)
			{
				throw invalid_ast_error("else blocks cannot be present in the final AST", else_block.location);
			}

			void apply(if_chain const & if_chain)
			{
				for (auto const & block : if_chain.blocks)
				{
					if (block.condition)
						apply(*block.condition);
					scopes.emplace_back();
					apply(*block.statements);
					scopes.pop_back();
				}
			}

			void apply(while_block const & while_block)
			{
				apply(*while_block.condition);
				scopes.emplace_back();
				apply(*while_block.statements);
				scopes.pop_back();
			}

			void apply(function_definition const & function_definition)
			{
				if (scopes.back().contains(function_definition.name))
					throw parse_error("Identifier \"" + function_definition.name + "\" is already defined at this scope", function_definition.location);

				std::unordered_set<std::string> argument_names;
				for (auto const & argument : function_definition.arguments)
				{
					if (argument_names.count(argument.name) > 0)
						throw parse_error("Duplicate argument name \"" + argument.name + "\" in function \"" + function_definition.name + "\"", argument.location);
					argument_names.insert(argument.name);
					apply(*argument.type);
				}

				apply(*function_definition.return_type);

				scopes.back().functions.insert(function_definition.name);

				auto & scope = scopes.emplace_back();
				scope.is_function_scope = true;
				scope.variables = std::move(argument_names);
				apply(*function_definition.statements);
				scopes.pop_back();
			}

			void apply(return_statement const & return_statement)
			{
				if (return_statement.value)
					apply(*return_statement.value);
			}

			void apply(field_definition const & field_definition)
			{
				apply(*field_definition.type);
			}

			void apply(struct_definition const & struct_definition)
			{
				if (scopes.back().contains(struct_definition.name))
					throw parse_error("Identifier \"" + struct_definition.name + "\" is already defined at this scope", struct_definition.location);

				for (auto const & field : struct_definition.fields)
					apply(field);

				scopes.back().structs.insert(struct_definition.name);
			}

		};

	}

	void resolve_identifiers(statement_list_ptr & statements)
	{
		resolve_identifiers_visitor visitor;
		visitor.scopes.emplace_back().is_global_scope = true;
		visitor.apply(*statements);
	}

}