#include <pslang/interpreter/exec.hpp>
#include <pslang/interpreter/eval.hpp>
#include <pslang/interpreter/error.hpp>
#include <pslang/ast/statement.hpp>
#include <pslang/types/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_.frame_stack.pop_back();
			}

		private:
			context & context_;
		};

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

		void exec_impl(context & context, ast::assignment const & assignment)
		{
			auto new_value = eval(context, assignment.rhs);
			auto new_type = type_of(new_value);

			auto ref = eval_ref(context, assignment.lhs);

			auto existing_type = type_of(*ref);
			if (!types::equal(existing_type, new_type))
			{
				std::ostringstream os;
				os << "Cannot assign a value of type ";
				types::print(os, new_type);
				os << " to a variable of type ";
				types::print(os, existing_type);
				throw internal_error(os.str(), assignment.location);
			}

			*ref = std::move(new_value);
		}

		void exec_impl(context & context, ast::variable_declaration const & variable_declaration)
		{
			auto & scope = context.frame_stack.back();
			if (scope.contains(variable_declaration.name))
				throw internal_error("Identifier \"" + variable_declaration.name + "\" is already defined in this scope", variable_declaration.location);

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

		void exec_impl(context & context, ast::if_block const & if_block)
		{
			throw internal_error("if blocks cannot be present in the final AST", if_block.location);
		}

		void exec_impl(context & context, ast::else_block const & else_block)
		{
			throw internal_error("else blocks cannot be present in the final AST", else_block.location);
		}

		void exec_impl(context & context, ast::else_if_block const & else_if_block)
		{
			throw internal_error("else if blocks cannot be present in the final AST", else_if_block.location);
		}

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

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

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

				if (!do_exec)
					continue;

				context.frame_stack.emplace_back();
				stack_pop_guard guard(context);
				exec(context, block.statements);
				break;
			}
		}

		void exec_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 (!types::equal(actual_type, types::primitive_type{types::bool_type{}}))
				{
					std::ostringstream os;
					os << "Expected type bool, got type ";
					types::print(os, actual_type);
					os << " in while block condition";
					throw internal_error(os.str(), get_location(*while_block.condition));
				}

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

		void exec_impl(context & context, ast::function_definition const & function_definition)
		{
			auto & frame = context.frame_stack.back();
			if (frame.contains(function_definition.name))
				throw std::runtime_error("Identifier \"" + function_definition.name + "\" is already defined in this scope");

			function_value value;

			for (auto const & argument : function_definition.arguments)
				value.arguments.push_back({.name = argument.name, .type = get_type(*argument.type)});
			value.return_type = get_type(*function_definition.return_type);
			value.statements = function_definition.statements;

			frame.variables[function_definition.name] = {.category = ast::value_category::constant, .value = std::move(value)};
		}

		void exec_impl(context & context, ast::foreign_function_declaration const & foreign_function_declaration)
		{
			auto & frame = context.frame_stack.back();
			if (frame.contains(foreign_function_declaration.name))
				throw std::runtime_error("Identifier \"" + foreign_function_declaration.name + "\" is already defined in this scope");

			foreign_function_value value;

			for (auto const & argument : foreign_function_declaration.arguments)
				value.arguments.push_back({.name = argument.name, .type = get_type(*argument.type)});
			value.return_type = get_type(*foreign_function_declaration.return_type);

			value.pointer = context.foreign_resolver(foreign_function_declaration.name);

			frame.variables[foreign_function_declaration.name] = {.category = ast::value_category::constant, .value = std::move(value)};
		}

		void exec_impl(context & context, ast::return_statement const & return_statement)
		{
			// NB: cannot use return_statement.level here because lexical scope stack
			// almost always differs from execution frame stack

			// NB: cannot use iterators here because they might get invalidated during eval()

			for (std::size_t n = context.frame_stack.size(); n --> 0;)
			{
				{
					auto & frame = context.frame_stack[n];

					if (!frame.expected_return_type)
						continue;
				}

				auto value = return_statement.value ? eval(context, return_statement.value) : unit_value{};
				auto actual_type = type_of(value);

				auto & frame = context.frame_stack[n];

				if (!types::equal(*frame.expected_return_type, actual_type))
				{
					std::ostringstream os;
					os << "Returning value of type ";
					types::print(os, actual_type);
					os << " from a function returning ";
					types::print(os, *frame.expected_return_type);
					throw internal_error(os.str(), return_statement.location);
				}

				frame.return_value = std::move(value);
				throw return_exception{};
			}

			throw internal_error("Cannot return outside of function scope", return_statement.location);
		}

		void exec_impl(context & context, ast::field_definition const & field_definition)
		{
			throw internal_error("Field definitions cannot be present in the final AST outside of struct definitions", field_definition.location);
		}

		void exec_impl(context & context, ast::struct_definition const & struct_definition)
		{
			auto & frame = context.frame_stack.back();
			if (frame.contains(struct_definition.name))
				throw internal_error("Identifier \"" + struct_definition.name + "\" is already defined in this scope", struct_definition.location);

			struct_data result;
			for (auto const & field : struct_definition.fields)
			{
				result.fields.push_back({
					.name = field.name,
					.type = get_type(*field.type),
				});
			}
			frame.structs[struct_definition.name] = std::move(result);
		}

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

	}

	void exec(context & context, ast::statement_list_ptr const & statements)
	{
		exec_impl(context, statements);
	}

}