#include <pslang/interpreter/interpreter.hpp>
#include <pslang/interpreter/eval.hpp>
#include <pslang/ast/statement.hpp>
#include <pslang/type/print.hpp>

#include <stdexcept>
#include <sstream>

namespace pslang::interpreter
{

	namespace
	{

		struct return_exception
		{};

		struct stack_pop_guard
		{
			stack_pop_guard(context & context)
				: context_(context)
			{}

			~stack_pop_guard()
			{
				context_.scope_stack.pop_back();
			}

		private:
			context & context_;
		};

		void execute_impl(context & context, ast::expression_ptr const & expression)
		{
			eval(context, expression);
		}

		void execute_impl(context & context, ast::assignment const & assignment)
		{
			std::string name;
			if (auto identifier = std::get_if<ast::identifier>(assignment.lhs.get()))
				name = identifier->name;
			else
				throw std::runtime_error("Cannot assign a value to a non-identifier");

			for (auto it = context.scope_stack.rbegin(); it != context.scope_stack.rend(); ++it)
			{
				if (auto jt = it->variables.find(name); jt != it->variables.end())
				{
					if (jt->second.category != ast::value_category::_mutable)
						throw std::runtime_error("Cannot assign a value to a non-mutable variable");

					auto new_value = eval(context, assignment.rhs);
					auto new_type = type_of(new_value);
					auto existing_type = type_of(jt->second.value);
					if (!type::equal(existing_type, new_type))
					{
						std::ostringstream os;
						os << "Cannot assign a value of type ";
						type::print(os, new_type);
						os << " to a variable of type ";
						type::print(os, existing_type);
						throw std::runtime_error(os.str());
					}
					jt->second.value = std::move(new_value);;
					return;
				}
			}

			throw std::runtime_error("Identifier \"" + name + "\" is not defined");
		}

		void execute_impl(context & context, ast::variable_declaration const & variable_declaration)
		{
			auto & scope = context.scope_stack.back();
			if (scope.variables.count(variable_declaration.name) > 0)
				throw std::runtime_error("Error: variable \"" + variable_declaration.name + "\" is already declared");

			auto value = eval(context, variable_declaration.initializer);
			if (variable_declaration.type)
			{
				auto actual_type = type_of(value);
				if (!type::equal(*variable_declaration.type, actual_type))
				{
					std::ostringstream os;
					os << "Cannot initialize a variable of type ";
					type::print(os, *variable_declaration.type);
					os << " with an expression of type ";
					type::print(os, actual_type);
					throw std::runtime_error(os.str());
				}
			}
			context.scope_stack.back().variables[variable_declaration.name] = {.category = variable_declaration.category, .value = value};
		}

		void execute_impl(context & context, ast::if_block const &)
		{
			throw std::runtime_error("Internal interpreter error: if blocks cannot be present in the final AST");
		}

		void execute_impl(context & context, ast::else_block const &)
		{
			throw std::runtime_error("Internal interpreter error: else blocks cannot be present in the final AST");
		}

		void execute_impl(context & context, ast::else_if_block const &)
		{
			throw std::runtime_error("Internal interpreter error: else if blocks cannot be present in the final AST");
		}

		void execute_impl(context & context, ast::if_chain const & if_chain)
		{
			for (auto const & block : if_chain.blocks)
			{
				bool do_execute = true;

				if (block.condition)
				{
					auto value = eval(context, block.condition);
					auto actual_type = type_of(value);
					if (!type::equal(actual_type, type::primitive_type{type::bool_type{}}))
					{
						std::ostringstream os;
						os << "Expected type bool, got type ";
						type::print(os, actual_type);
						os << " in if block condition";
						throw std::runtime_error(os.str());
					}

					do_execute = std::get<bool_value>(std::get<primitive_value>(value)).value;
				}

				if (!do_execute)
					continue;

				context.scope_stack.emplace_back();
				stack_pop_guard guard(context);
				execute(context, block.statements);
				break;
			}
		}

		void execute_impl(context & context, ast::while_block const & while_block)
		{
			while (true)
			{
				auto value = eval(context, while_block.condition);
				auto actual_type = type_of(value);
				if (!type::equal(actual_type, type::primitive_type{type::bool_type{}}))
				{
					std::ostringstream os;
					os << "Expected type bool, got type ";
					type::print(os, actual_type);
					os << " in while block condition";
					throw std::runtime_error(os.str());
				}

				if (std::get<bool_value>(std::get<primitive_value>(value)).value)
				{
					context.scope_stack.emplace_back();
					stack_pop_guard guard(context);
					execute(context, while_block.statements);
				}
				else
					break;
			}
		}

		void execute_impl(context & context, ast::function_definition const & function_definition)
		{
			auto & scope = context.scope_stack.back();
			if (scope.functions.count(function_definition.name) > 0)
				throw std::runtime_error("Function \"" + function_definition.name + "\" is already defined");

			auto & function = scope.functions[function_definition.name];
			function.arguments = function_definition.arguments;
			function.return_type = function_definition.return_type;
			function.statements = function_definition.statements;
		}

		void execute_impl(context & context, ast::return_statement const & return_statement)
		{
			auto value = return_statement.value ? eval(context, return_statement.value) : unit_value{};
			for (auto it = context.scope_stack.rbegin(); it != context.scope_stack.rend(); ++it)
			{
				if (it->is_function_scope)
				{
					it->return_value = std::move(value);
					throw return_exception{};
				}
			}

			throw std::runtime_error("Cannot return outside of a function");
		}

		void execute_impl(context & context, ast::statement_list_ptr const & statements)
		{
			for (auto const & statement : statements->statements)
			{
				try
				{
					std::visit([&](auto const & statement){ execute_impl(context, statement); }, *statement);
				}
				catch (return_exception const &)
				{
					if (context.scope_stack.back().is_function_scope)
						break;
					else
						throw;
				}
			}
		}

	}

	void execute(context & context, ast::statement_list_ptr const & statements)
	{
		execute_impl(context, statements);
	}

}