#include <pslang/ast/print.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/types/print.hpp>
#include <pslang/types/type.hpp>

#include <iomanip>

namespace pslang::ast
{

	namespace
	{

		print_options as_child(print_options options)
		{
			options.indent_level += 1;
			return options;
		}

		void put_indent(std::ostream & out, print_options const & options)
		{
			for (std::size_t i = 0; i < options.indent_level; ++i)
				out << options.indent_string;
		}

		struct type_print_visitor
		{
			std::ostream & out;

			void operator()(types::unit_type const &)
			{
				out << "unit";
			}

			template <typename T>
			void operator()(types::primitive_type_base<T> const & type)
			{
				types::print(out, types::primitive_type{type});
			}

			template <typename Self>
			void operator()(Self & self, array_type const & type)
			{
				self(*type.element_type);
				out << "[" << type.size << "]";
			}

			template <typename Self>
			void operator()(Self & self, function_type const & type)
			{
				if (type.arguments.size() == 1)
				{
					self(*type.arguments.front());
					out << " -> ";
					self(*type.result);
					return;
				}

				out << '(';
				bool first = true;
				for (auto const & argument : type.arguments)
				{
					if (!first) out << ", ";
					first = false;
					self(*argument);
				}
				out << ") -> ";
				self(*type.result);
			}

			void operator()(type_identifier const & type)
			{
				out << type.name;
			}
		};

		struct expression_print_visitor
		{
			std::ostream & out;
			print_options options;

			template <typename Self, typename Node>
			void child(Self & self, Node const & node)
			{
				++options.indent_level;
				self(node);
				--options.indent_level;
			}

			void operator()(bool_literal const & node)
			{
				put_indent(out, options);
				out << "bool literal { value = " << (node.value ? "true" : "false") << " }\n";
			}

			void operator()(i8_literal const & node)
			{
				put_indent(out, options);
				out << "i8 literal { value = " << (std::int32_t)node.value << " }\n";
			}

			void operator()(u8_literal const & node)
			{
				put_indent(out, options);
				out << "u8 literal { value = " << (std::uint32_t)node.value << " }\n";
			}

			void operator()(i16_literal const & node)
			{
				put_indent(out, options);
				out << "i16 literal { value = " << node.value << " }\n";
			}

			void operator()(u16_literal const & node)
			{
				put_indent(out, options);
				out << "u16 literal { value = " << node.value << " }\n";
			}

			void operator()(i32_literal const & node)
			{
				put_indent(out, options);
				out << "i32 literal { value = " << node.value << " }\n";
			}

			void operator()(u32_literal const & node)
			{
				put_indent(out, options);
				out << "u32 literal { value = " << node.value << " }\n";
			}

			void operator()(i64_literal const & node)
			{
				put_indent(out, options);
				out << "i64 literal { value = " << node.value << " }\n";
			}

			void operator()(u64_literal const & node)
			{
				put_indent(out, options);
				out << "u64 literal { value = " << node.value << " }\n";
			}

			void operator()(f32_literal const & node)
			{
				put_indent(out, options);
				out << "f32 literal { value = " << std::setprecision(7) << node.value << " }\n";
			}

			void operator()(f64_literal const & node)
			{
				put_indent(out, options);
				out << "f64 literal { value = " << std::setprecision(15) << node.value << " }\n";
			}

			void operator()(identifier const & node)
			{
				put_indent(out, options);
				out << "identifier { name = \"" << node.name << "\" }\n";
			}

			template <typename Self>
			void operator()(Self & self, unary_operation const & node)
			{
				put_indent(out, options);
				out << node.type << '\n';
				child(self, *node.arg1);
			}

			template <typename Self>
			void operator()(Self & self, binary_operation const & node)
			{
				put_indent(out, options);
				out << node.type << '\n';
				child(self, *node.arg1);
				child(self, *node.arg2);
			}

			template <typename Self>
			void operator()(Self & self, cast_operation const & node)
			{
				put_indent(out, options);
				out << "cast as ";
				print(out, *node.type);
				out << '\n';
				child(self, *node.expression);
			}

			template <typename Self>
			void operator()(Self & self, function_call const & node)
			{
				put_indent(out, options);
				out << "call\n";
				child(self, *node.function);
				for (auto const & argument : node.arguments)
					child(self, *argument);
			}

			template <typename Self>
			void operator()(Self & self, array const & node)
			{
				put_indent(out, options);
				out << "array\n";
				for (auto const & element : node.elements)
					child(self, *element);
			}

			template <typename Self>
			void operator()(Self & self, array_access const & node)
			{
				put_indent(out, options);
				out << "array access\n";
				child(self, *node.array);
				child(self, *node.index);
			}

			template <typename Self>
			void operator()(Self & self, field_access const & node)
			{
				put_indent(out, options);
				out << "field access { name = \"" << node.field_name << "\" }\n";
				child(self, *node.object);
			}
		};

		struct statement_print_visitor
		{
			std::ostream & out;
			print_options options;

			template <typename Self, typename Node>
			void child(Self & self, Node const & node)
			{
				++options.indent_level;
				self(node);
				--options.indent_level;
			}

			void operator()(expression_ptr const & node)
			{
				print(out, *node, options);
			}

			template <typename Self>
			void operator()(Self & self, assignment const & node)
			{
				put_indent(out, options);
				out << "assignment\n";
				child(self, node.lhs);
				child(self, node.rhs);
			}

			template <typename Self>
			void operator()(Self & self, variable_declaration const & node)
			{
				put_indent(out, options);
				out << "variable declaration { category = " << node.category << ", name = \"" << node.name << "\"";
				if (node.type)
				{
					out << ", type = ";
					print(out, *node.type);
				}
				out << " }\n";
				child(self, node.initializer);
			}

			template <typename Self>
			void operator()(Self & self, if_block const & node)
			{
				put_indent(out, options);
				out << "if\n";
				child(self, node.condition);
			}

			void operator()(else_block const & node)
			{
				put_indent(out, options);
				out << "else\n";
			}

			template <typename Self>
			void operator()(Self & self, else_if_block const & node)
			{
				put_indent(out, options);
				out << "else if\n";
				child(self, node.condition);
			}

			template <typename Self>
			void operator()(Self & self, if_chain const & node)
			{
				put_indent(out, options);
				out << "if chain\n";
				++options.indent_level;
				for (auto const & block : node.blocks)
				{
					put_indent(out, options);
					out << "condition\n";
					if (block.condition)
						child(self, block.condition);
					else
					{
						put_indent(out, as_child(options));
						out << "(none)\n";
					}

					put_indent(out, options);
					out << "body\n";
					child(self, *block.statements);
				}
				--options.indent_level;
			}

			template <typename Self>
			void operator()(Self & self, while_block const & node)
			{
				put_indent(out, options);
				out << "while\n";
				child(self, node.condition);
				child(self, *node.statements);
			}

			template <typename Self>
			void operator()(Self & self, function_definition const & node)
			{
				put_indent(out, options);
				out << "function { name = \"" << node.name << "\", return type = ";
				print(out, *node.return_type);
				out << " }\n";
				++options.indent_level;
				for (auto const & arg : node.arguments)
				{
					put_indent(out, options);
					out << "argument { name = \"" << arg.name << "\", type = ";
					print(out, *arg.type);
					out << " }\n";
				}
				put_indent(out, options);
				out << "body\n";
				child(self, *node.statements);
			}

			template <typename Self>
			void operator()(Self & self, return_statement const & node)
			{
				put_indent(out, options);
				out << "return\n";
				if (node.value)
					child(self, node.value);
			}

			void operator()(field_definition const & node)
			{
				put_indent(out, options);
				out << "field { name = \"" << node.name << "\", type = ";
				print(out, *node.type);
				out << " }\n";
			}

			template <typename Self>
			void operator()(Self & self, struct_definition const & node)
			{
				put_indent(out, options);
				out << "struct { name = \"" << node.name << "\" }\n";
				for (auto const & field : node.fields)
					child(self, field);
			}
		};

	}

	void print(std::ostream & out, type const & node)
	{
		apply(type_print_visitor{out}, node);
	}

	void print(std::ostream & out, expression const & node, print_options const & options)
	{
		apply(expression_print_visitor{out, options}, node);
	}

	void print(std::ostream & out, statement const & node, print_options const & options)
	{
		apply(statement_print_visitor{out, options}, node);
	}

	void print(std::ostream & out, statement_list const & node, print_options const & options)
	{
		apply(statement_print_visitor{out, options}, node);
	}

}