#include <pslang/ast/print.hpp>
#include <pslang/type/print.hpp>

#include <iomanip>

namespace pslang::ast
{

	namespace
	{

		print_options 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;
		}

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

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

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

		void print_impl(std::ostream & out, i16_literal const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "i16 literal { value = " << node.value << " }\n";
		}

		void print_impl(std::ostream & out, u16_literal const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "u16 literal { value = " << node.value << " }\n";
		}

		void print_impl(std::ostream & out, i32_literal const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "i32 literal { value = " << node.value << " }\n";
		}

		void print_impl(std::ostream & out, u32_literal const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "u32 literal { value = " << node.value << " }\n";
		}

		void print_impl(std::ostream & out, i64_literal const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "i64 literal { value = " << node.value << " }\n";
		}

		void print_impl(std::ostream & out, u64_literal const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "u64 literal { value = " << node.value << " }\n";
		}

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

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

		void print_impl(std::ostream & out, literal const & node, print_options const & options)
		{
			std::visit([&](auto const & value){ print_impl(out, value, options); }, node);
		}

		void print_impl(std::ostream & out, identifier const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "identifier { name = \"" << node.name << "\" }\n";
		}

		void print_impl(std::ostream & out, unary_operation const & node, print_options const & options)
		{
			put_indent(out, options);
			out << node.type << '\n';
			print(out, node.arg1, child(options));
		}

		void print_impl(std::ostream & out, binary_operation const & node, print_options const & options)
		{
			put_indent(out, options);
			out << node.type << '\n';
			print(out, node.arg1, child(options));
			print(out, node.arg2, child(options));
		}

		void print_impl(std::ostream & out, cast_operation const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "cast as ";
			type::print(out, *node.type);
			out << '\n';
			print(out, node.expression, child(options));
		}

		void print_impl(std::ostream & out, function_call const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "call\n";
			print(out, node.function, child(options));
			for (auto const & argument : node.arguments)
				print(out, argument, child(options));
		}

		void print_impl(std::ostream & out, array const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "array\n";
			for (auto const & element : node.elements)
				print(out, element, child(options));
		}

		void print_impl(std::ostream & out, array_access const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "array access\n";
			print(out, node.array, child(options));
			print(out, node.index, child(options));
		}

		void print_impl(std::ostream & out, field_access const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "field access { name = \"" << node.field_name << "\" }\n";
			print(out, node.object, child(options));
		}

		void print_impl(std::ostream & out, expression_ptr const & node, print_options const & options)
		{
			std::visit([&](auto const & value){ print_impl(out, value, options); }, *node);
		}

		void print_impl(std::ostream & out, assignment const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "assignment\n";
			print(out, node.lhs, child(options));
			print(out, node.rhs, child(options));
		}

		void print_impl(std::ostream & out, variable_declaration const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "variable declaration { category = " << node.category << ", name = \"" << node.name << "\"";
			if (node.type)
			{
				out << ", type = ";
				type::print(out, *node.type);
			}
			out << " }\n";
			print(out, node.initializer, child(options));
		}

		void print_impl(std::ostream & out, if_block const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "if\n";
			print(out, node.condition, child(options));
		}

		void print_impl(std::ostream & out, else_block const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "else\n";
		}

		void print_impl(std::ostream & out, else_if_block const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "else if\n";
			print(out, node.condition, child(options));
		}

		void print_impl(std::ostream & out, if_chain const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "if chain\n";
			for (auto const & block : node.blocks)
			{
				put_indent(out, child(options));
				out << "condition\n";
				if (block.condition)
					print(out, block.condition, child(child(options)));
				else
				{
					put_indent(out, child(child(options)));
					out << "(none)\n";
				}

				put_indent(out, child(options));
				out << "body\n";
				print(out, block.statements, child(child(options)));
			}
		}

		void print_impl(std::ostream & out, while_block const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "while\n";
			print(out, node.condition, child(options));
			print(out, node.statements, child(options));
		}

		void print_impl(std::ostream & out, function_definition const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "function { name = \"" << node.name << "\", return type = ";
			type::print(out, *node.return_type);
			out << " }\n";
			for (auto const & arg : node.arguments)
			{
				put_indent(out, child(options));
				out << "argument { name = \"" << arg.name << "\", type = ";
				type::print(out, *arg.type);
				out << " }\n";
			}
			put_indent(out, child(options));
			out << "body\n";
			print(out, node.statements, child(child(options)));
		}

		void print_impl(std::ostream & out, return_statement const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "return\n";
			if (node.value)
				print(out, node.value, child(options));
		}

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

		void print_impl(std::ostream & out, struct_definition const & node, print_options const & options)
		{
			put_indent(out, options);
			out << "struct { name = \"" << node.name << "\" }\n";
			for (auto const & field : node.fields)
				print_impl(out, field, child(options));
		}

		void print_impl(std::ostream & out, statement_ptr const & node, print_options const & options)
		{
			std::visit([&](auto const & value){ print_impl(out, value, options); }, *node);
		}

		void print_impl(std::ostream & out, statement_list_ptr const & node, print_options const & options)
		{
			for (auto const & statement : node->statements)
				print(out, statement, options);
		}

	}

	void print(std::ostream & out, expression_ptr const & node, print_options const & options)
	{
		print_impl(out, node, options);
	}

	void print(std::ostream & out, statement_ptr const & node, print_options const & options)
	{
		print_impl(out, node, options);
	}

	void print(std::ostream & out, statement_list_ptr const & node, print_options const & options)
	{
		print_impl(out, node, options);
	}
}