#include <pslang/parser/parser.hpp>
#include <pslang/parser/error.hpp>
#include <pslang/interpreter/exec.hpp>
#include <pslang/interpreter/error.hpp>
#include <pslang/ast/statement.hpp>
#include <pslang/ast/preprocess.hpp>
#include <pslang/ast/print.hpp>
#include <pslang/ir/compiler.hpp>
#include <pslang/ir/print.hpp>
#include <pslang/jit/jit.hpp>
#include <pslang/jit/executable.hpp>
#include <pslang/jit/foreign.hpp>

#include <filesystem>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cmath>

std::string extract_nth_line(std::filesystem::path const & path, std::size_t n)
{
	std::ifstream file(path);
	std::string line;

	for (size_t i = 1; i <= n && std::getline(file, line); ++i) {
		if (i == n) {
			return line;
		}
	}

	return {};
}

std::size_t replace_tabs_with_spaces(std::string & str, std::size_t count)
{
	std::string replaced;
	for (char c : str)
	{
		if (c == '\t')
			replaced.append(count, ' ');
		else
			replaced.append(1, c);
	}
	std::swap(str, replaced);
	return str.size() - replaced.size();
}

void print_error_context(std::filesystem::path const & file, pslang::ast::location const & location)
{
	std::size_t max_line_number_digits = std::floor(std::log10(location.end.line)) + 1;

	if (location.begin.line == location.end.line)
	{
		auto line = extract_nth_line(file, location.begin.line);
		auto extra = replace_tabs_with_spaces(line, 2);

		std::cerr << std::endl;
		std::cerr << "line " << std::setw(max_line_number_digits) << std::right << location.begin.line << ": ";
		std::cerr << line << std::endl;
		std::cerr << std::string(std::max<std::size_t>(location.begin.column, 1) - 1 + max_line_number_digits + 7 + extra, ' ') << std::string(location.end.column - location.begin.column, '^') << std::endl;
	}
	else
	{
		std::cerr << std::endl;
		for (std::size_t line = location.begin.line; line <= location.end.line; ++line)
		{
			std::cerr << "line " << std::setw(max_line_number_digits) << std::right << line << ": ";
			std::cerr << extract_nth_line(file, line) << std::endl;
		}
	}
}

int main(int argc, char ** argv)
{
	if (argc == 1)
	{
		std::cout << "Usage: psli [ options ] <file1> [ <file2> ... ]\n";
		std::cout << "Available options:\n";
		std::cout << "    -t, --trace        Trace each line of execution\n";
		std::cout << "    -d, --dump         Dump all variables after processing all files\n";
		std::cout << "    -p, --print        Print the AST after parsing each file\n";
		std::cout << "    -i, --ir           Print the IR after parsing each file\n";
		std::cout << "    -j, --jit          Just-in-time compile & execute the code instead of interpreting\n";
		return 0;
	}

	using namespace pslang;

	auto context = interpreter::empty_context();
	context.foreign_resolver = &jit::load_foreign;

	bool dump = false;
	bool dump_ast = false;
	bool dump_ir = false;
	bool jit = false;

	std::vector<std::string> filenames;
	std::vector<ast::statement_ptr> parsed;
	std::vector<ir::module_context> ir_compiled;

	bool no_more_options = false;

	for (int arg = 1; arg < argc; ++arg)
	{
		if (argv[arg][0] == 0)
			continue;

		if (std::strcmp(argv[arg], "--") == 0)
		{
			no_more_options = true;
			continue;
		}

		if (!no_more_options)
		{
			if (argv[arg][0] == '-' && argv[arg][1] != 0 && argv[arg][1] != '-')
			{
				for (char * p = argv[arg] + 1; *p != 0; ++p)
				{
					if (*p == 'd')
						dump = true;
					else if (*p == 't')
						context.trace = true;
					else if (*p == 'p')
						dump_ast = true;
					else if (*p == 'i')
						dump_ir = true;
					else if (*p == 'j')
						jit = true;
					else
					{
						std::cerr << "Unknown option '" << *p << "'\n";
						return EXIT_FAILURE;
					}
				}
				continue;
			}

			if (std::strcmp(argv[arg], "--dump") == 0)
			{
				dump = true;
				continue;
			}

			if (std::strcmp(argv[arg], "--trace") == 0)
			{
				context.trace = true;
				continue;
			}

			if (std::strcmp(argv[arg], "--print") == 0)
			{
				dump_ast = true;
				continue;
			}

			if (std::strcmp(argv[arg], "--ir") == 0)
			{
				dump_ir = true;
				continue;
			}

			if (std::strcmp(argv[arg], "--jit") == 0)
			{
				jit = true;
				continue;
			}
		}

		if (!std::filesystem::exists(argv[arg]))
		{
			std::cerr << "Error: input file \"" << argv[arg] << "\" does not exist\n";
			return EXIT_FAILURE;
		}

		if (std::filesystem::is_directory(argv[arg]))
		{
			std::cerr << "Error: input file \"" << argv[arg] << "\" is a directory\n";
			return EXIT_FAILURE;
		}

		try
		{
			filenames.push_back(argv[arg]);
			auto root = parser::parse(filenames.back());

			if (dump_ast)
			{
				std::cout << "Input file " << filenames.back() << " AST dump:\n\n";
				if (auto function_definition = std::get_if<ast::function_definition>(parsed.back().get()))
					ast::print(std::cout, *function_definition->statements);
				std::cout << "\n" << std::flush;
			}

			ast::resolve_identifiers(root);
			ast::check_and_infer_types(root);
			ast::validate(root);
			parsed.push_back(std::move(root));

			ir_compiled.emplace_back();
			ir::compile(ir_compiled.back(), parsed.back());

			if (dump_ir)
			{
				std::cout << "Input file " << filenames.back() << " IR dump:\n\n";
				ir::print(std::cout, ir_compiled.back());
				std::cout << "\n" << std::flush;
			}
		}
		catch (ast::parse_error const & error)
		{
			std::cerr << "Parse error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (ast::type_error const & error)
		{
			std::cerr << "Type error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (ast::invalid_ast_error const & error)
		{
			std::cerr << "Invalid AST at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (ast::validation_error const & error)
		{
			std::cerr << "Validation error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (interpreter::internal_error const & error)
		{
			std::cerr << "Internal error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
		catch (interpreter::runtime_error const & error)
		{
			std::cerr << "Runtime error at " << error.location() << ":\n    " << error.what() << std::endl;
			print_error_context(argv[arg], error.location());
			return EXIT_FAILURE;
		}
	}

	if (jit)
	{
		// TODO: treat all input files as modules combined into a single program
		for (std::size_t i = 0; i < filenames.size(); ++i)
		{
			jit::program_context pcontext
			{
				.abi = jit::host_abi(),
			};

			jit::compile(pcontext, ir_compiled[i]);

			for (auto const & resolve : pcontext.foreign_resolve)
			{
				auto fptr = jit::load_foreign(resolve.name);
				std::copy_n((std::uint8_t const *)(&fptr), 8, pcontext.storage.storage.data() + resolve.offset);
			}

			auto executable = jit::make_host_executable(pcontext.storage);

			auto entry_point = (void(*)())(executable.get() + pcontext.entry_point);
			entry_point();
		}
	}
	else
	{
		// for (auto const & ast : parsed)
		// 	interpreter::exec(context, ast);

		if (dump)
			interpreter::dump(std::cout, context);
	}
}